// // Copyright 2026 The InfiniFlow Authors. All Rights Reserved. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. // package elasticsearch import ( "bytes" "cmp" "context" "encoding/json" "fmt" "io" "net/http" "os" "reflect" "regexp" "slices" "sort" "strconv" "strings" "github.com/elastic/go-elasticsearch/v8/esapi" "github.com/json-iterator/go" "ragflow/internal/common" "ragflow/internal/engine/types" "go.uber.org/zap" ) var jsonIterator = jsoniter.Config{ SortMapKeys: false, }.Froze() var ( elasticsearchHighlightEmTagRE = regexp.MustCompile(`[^<>]+`) elasticsearchHighlightNewlineRE = regexp.MustCompile(`[\r\n]`) elasticsearchHighlightDelimiterRE = regexp.MustCompile(`[.?!;\n]`) elasticsearchLetterRE = regexp.MustCompile(`\pL`) elasticsearchEnglishLetterRE = regexp.MustCompile(`[A-Za-z]`) ) // CreateChunkStore creates an index func (e *elasticsearchEngine) CreateChunkStore(ctx context.Context, baseName, datasetID string, vectorSize int, parserID string) error { if baseName == "" { return fmt.Errorf("index name cannot be empty") } // Check if index already exists exists, err := e.indexExists(ctx, baseName) if err != nil { return fmt.Errorf("failed to check index existence: %w", err) } if exists { common.Info("Index already exists, skipping creation", zap.String("index_name", baseName)) return nil } // Load mapping based on index type var mapping map[string]interface{} if datasetID == "skill" { // Load skill-specific mapping skillMapping, err := loadSkillMapping() if err != nil { return fmt.Errorf("failed to load skill mapping: %w", err) } mapping = skillMapping } else { // Default mapping for dataset mapping = map[string]interface{}{ "settings": map[string]interface{}{ "number_of_shards": 1, "number_of_replicas": 0, }, } } // Prepare request body var body io.Reader if mapping != nil { data, err := json.Marshal(mapping) if err != nil { return fmt.Errorf("failed to marshal mapping: %w", err) } body = bytes.NewReader(data) } // Create index req := esapi.IndicesCreateRequest{ Index: baseName, Body: body, } res, err := req.Do(ctx, e.client) if err != nil { return fmt.Errorf("failed to create index: %w", err) } defer res.Body.Close() if res.IsError() { bodyBytes, _ := io.ReadAll(res.Body) reason := extractErrorReason(bodyBytes) if reason != "" { return fmt.Errorf("elasticsearch error: %s", reason) } return fmt.Errorf("elasticsearch returned error: %s, body: %s", res.Status(), string(bodyBytes)) } // Parse response var result map[string]interface{} if err := json.NewDecoder(res.Body).Decode(&result); err != nil { return fmt.Errorf("failed to parse response: %w", err) } acknowledged, ok := result["acknowledged"].(bool) if !ok || !acknowledged { return fmt.Errorf("index creation not acknowledged") } common.Info("Successfully created Elasticsearch index", zap.String("index_name", baseName)) return nil } // InsertChunks inserts chunks into a chunk index // If a chunk with the same id + doc_id + kb_id already exists, it will be updated with the new value func (e *elasticsearchEngine) InsertChunks(ctx context.Context, chunks []map[string]interface{}, baseName string, datasetID string) ([]string, error) { common.Info("ElasticsearchConnection.InsertChunks called", zap.String("index_name", baseName), zap.Int("chunkCount", len(chunks))) if len(chunks) == 0 { return []string{}, nil } if baseName == "" { return nil, fmt.Errorf("index name cannot be empty") } // Build bulk request body with index operations (upsert behavior: insert if not exists, update if exists) var buf bytes.Buffer for _, doc := range chunks { docID, _ := doc["doc_id"].(string) chunkID, _ := doc["id"].(string) if docID == "" || chunkID == "" { common.Warn("Skipping chunk without doc_id or id") continue } // Action line: use json.Marshal to properly escape string values action, err := json.Marshal(map[string]interface{}{ "index": map[string]interface{}{ "_index": baseName, "_id": chunkID, }, }) if err != nil { common.Error("Failed to marshal bulk action", err) return nil, fmt.Errorf("failed to marshal bulk action: %w", err) } buf.Write(action) buf.WriteByte('\n') // Document line: work with a copy to avoid mutating the original docCopy := copyFields(doc) docCopy["kb_id"] = datasetID if err := jsonIterator.NewEncoder(&buf).Encode(docCopy); err != nil { return nil, fmt.Errorf("failed to encode document: %w", err) } } // Execute bulk request with refresh="wait_for" req := esapi.BulkRequest{ Body: bytes.NewReader(buf.Bytes()), Refresh: "wait_for", } res, err := req.Do(ctx, e.client) if err != nil { common.Error("Failed to execute bulk request", err) return nil, fmt.Errorf("failed to execute bulk request: %w", err) } defer res.Body.Close() if res.IsError() { bodyBytes, _ := io.ReadAll(res.Body) common.Sugar.Errorw("Elasticsearch bulk request returned error", "status", res.Status(), "body", string(bodyBytes)) return nil, fmt.Errorf("elasticsearch bulk request returned error: %s, body: %s", res.Status(), string(bodyBytes)) } // Parse bulk response var bulkResponse map[string]interface{} if err := json.NewDecoder(res.Body).Decode(&bulkResponse); err != nil { common.Error("Failed to parse bulk response", err) return nil, fmt.Errorf("failed to parse bulk response: %w", err) } // Check for errors in bulk response if errors, ok := bulkResponse["errors"].(bool); ok && errors { common.Warn("Bulk request had some errors") // Could iterate through items to find specific errors if needed } common.Info("ElasticsearchConnection.InsertChunks result", zap.String("index_name", baseName), zap.Int("count", len(chunks))) return []string{}, nil } // UpdateChunks updates chunks by condition func (e *elasticsearchEngine) UpdateChunks(ctx context.Context, condition map[string]interface{}, newValue map[string]interface{}, baseName string, datasetID string) error { fullIndexName := baseName common.Info("ElasticsearchConnection.UpdateChunks called", zap.String("index_name", fullIndexName), zap.Any("condition", condition), zap.Any("new_value", newValue)) if fullIndexName == "" { return fmt.Errorf("index name cannot be empty") } // Check if index exists exists, err := e.indexExists(ctx, fullIndexName) if err != nil { common.Error("Failed to check index existence", err) return fmt.Errorf("failed to check index existence: %w", err) } if !exists { return fmt.Errorf("index '%s' does not exist", fullIndexName) } // Add kb_id to condition condition["kb_id"] = datasetID // Case 1: Single document update (when condition["id"] is a string) if chunkID, ok := condition["id"].(string); ok { return e.updateSingleChunk(ctx, fullIndexName, chunkID, newValue) } // Case 2: Multi-document update via UpdateByQuery return e.updateChunksByQuery(ctx, fullIndexName, condition, newValue) } // updateSingleChunk handles single document update func (e *elasticsearchEngine) updateSingleChunk(ctx context.Context, indexName, chunkID string, newValue map[string]interface{}) error { common.Debug("ElasticsearchConnection.updateSingleChunk called", zap.String("indexName", indexName), zap.String("chunkID", chunkID)) // First find the document by id field to get the actual _id searchReq := map[string]interface{}{ "query": map[string]interface{}{ "term": map[string]interface{}{"id": chunkID}, }, } body, err := json.Marshal(searchReq) if err != nil { return fmt.Errorf("failed to marshal search request: %w", err) } res, err := e.client.Search( e.client.Search.WithContext(ctx), e.client.Search.WithIndex(indexName), e.client.Search.WithBody(bytes.NewReader(body)), ) if err != nil { return fmt.Errorf("failed to search for chunk: %w", err) } defer res.Body.Close() if res.IsError() { return fmt.Errorf("failed to search for chunk: %s", res.Status()) } var searchResult map[string]interface{} if err := json.NewDecoder(res.Body).Decode(&searchResult); err != nil { return fmt.Errorf("failed to parse search response: %w", err) } hits, ok := searchResult["hits"].(map[string]interface{}) if !ok { return fmt.Errorf("%w: %s", types.ErrDocumentNotFound, chunkID) } hitList, ok := hits["hits"].([]interface{}) if !ok || len(hitList) == 0 { return fmt.Errorf("%w: %s", types.ErrDocumentNotFound, chunkID) } firstHit, ok := hitList[0].(map[string]interface{}) if !ok { return fmt.Errorf("%w: %s", types.ErrDocumentNotFound, chunkID) } actualID, ok := firstHit["_id"].(string) if !ok { return fmt.Errorf("%w: %s", types.ErrDocumentNotFound, chunkID) } doc := copyFields(newValue) delete(doc, "id") removeValue, _ := doc["remove"] delete(doc, "remove") removeField, _ := removeValue.(string) removeDict, _ := removeValue.(map[string]interface{}) // Remove *_feas fields var feasFields []string for k := range doc { if strings.HasSuffix(k, "feas") { feasFields = append(feasFields, k) } } for _, k := range feasFields { scriptBody := map[string]interface{}{ "script": map[string]interface{}{ "source": fmt.Sprintf("ctx._source.remove(\"%s\");", k), }, } body, _ := json.Marshal(scriptBody) req := esapi.UpdateRequest{ Index: indexName, DocumentID: actualID, Body: bytes.NewReader(body), } res, err := req.Do(ctx, e.client) if err != nil { common.Warn("Failed to remove feas field", zap.String("field", k), zap.Error(err)) } else { res.Body.Close() } } // Remove specific field if removeField is set if removeField != "" { scriptBody := map[string]interface{}{ "script": map[string]interface{}{ "source": fmt.Sprintf("ctx._source.remove('%s');", removeField), }, } body, _ := json.Marshal(scriptBody) req := esapi.UpdateRequest{ Index: indexName, DocumentID: actualID, Body: bytes.NewReader(body), } res, err := req.Do(ctx, e.client) if err != nil { common.Warn("Failed to remove field", zap.String("field", removeField), zap.Error(err)) } else { res.Body.Close() } } // Remove specific values from array fields (removeDict) if removeDict != nil { scripts := []string{} params := make(map[string]interface{}) for kk, vv := range removeDict { scripts = append(scripts, fmt.Sprintf("if (ctx._source.containsKey('%s') && ctx._source.%s != null) { int i = ctx._source.%s.indexOf(params.p_%s); if (i >= 0) { ctx._source.%s.remove(i); }}", kk, kk, kk, kk, kk)) params[fmt.Sprintf("p_%s", kk)] = vv } if scripts != nil { scriptBody := map[string]interface{}{ "script": map[string]interface{}{ "source": strings.Join(scripts, ""), "params": params, }, } body, _ := json.Marshal(scriptBody) req := esapi.UpdateRequest{ Index: indexName, DocumentID: actualID, Body: bytes.NewReader(body), } res, err := req.Do(ctx, e.client) if err != nil { common.Warn("Failed to remove dict fields", zap.Error(err)) } else { res.Body.Close() } } } // Update document fields if any remain if len(doc) > 0 { updateBody := map[string]interface{}{"doc": doc} body, _ := json.Marshal(updateBody) req := esapi.UpdateRequest{ Index: indexName, DocumentID: actualID, Body: bytes.NewReader(body), } res, err := req.Do(ctx, e.client) if err != nil { return fmt.Errorf("failed to update document: %w", err) } defer res.Body.Close() if res.IsError() { if res.StatusCode == http.StatusNotFound { return fmt.Errorf("%w: %s", types.ErrDocumentNotFound, chunkID) } return fmt.Errorf("elasticsearch update error: %s", res.Status()) } } common.Debug("ElasticsearchConnection.updateSingleChunk completed", zap.String("indexName", indexName), zap.String("chunkID", chunkID)) return nil } // updateChunksByQuery handles multi-document update func (e *elasticsearchEngine) updateChunksByQuery(ctx context.Context, indexName string, condition map[string]interface{}, newValue map[string]interface{}) error { common.Debug("ElasticsearchConnection.updateChunksByQuery called", zap.String("indexName", indexName)) // Build bool query from condition var mustClauses []map[string]interface{} for k, v := range condition { if k == "exists" { mustClauses = append(mustClauses, map[string]interface{}{ "exists": map[string]interface{}{"field": v}, }) continue } if v == nil || v == "" { continue } if listVal, ok := v.([]interface{}); ok { mustClauses = append(mustClauses, map[string]interface{}{ "terms": map[string]interface{}{k: listVal}, }) } else if _, ok := v.(string); ok { mustClauses = append(mustClauses, map[string]interface{}{ "term": map[string]interface{}{k: v}, }) } else if _, ok := v.(int); ok { mustClauses = append(mustClauses, map[string]interface{}{ "term": map[string]interface{}{k: v}, }) } } boolQuery := map[string]interface{}{ "bool": map[string]interface{}{ "filter": mustClauses, }, } // Build painless scripts from newValue var scripts []string params := make(map[string]interface{}) for k, v := range newValue { if k == "remove" { if removeStr, ok := v.(string); ok { scripts = append(scripts, fmt.Sprintf("ctx._source.remove('%s');", removeStr)) continue } if removeDict, ok := v.(map[string]interface{}); ok { for kk, vv := range removeDict { scripts = append(scripts, fmt.Sprintf("if (ctx._source.containsKey('%s') && ctx._source.%s != null) { int i = ctx._source.%s.indexOf(params.p_%s); if (i >= 0) { ctx._source.%s.remove(i); }}", kk, kk, kk, kk, kk)) params[fmt.Sprintf("p_%s", kk)] = vv } } continue } if k == "add" { if addDict, ok := v.(map[string]interface{}); ok { for kk, vv := range addDict { vvStr, ok := vv.(string) if ok { vvStr = strings.TrimSpace(vvStr) scripts = append(scripts, fmt.Sprintf("ctx._source.%s.add(params.pp_%s);", kk, kk)) params[fmt.Sprintf("pp_%s", kk)] = vvStr } } } continue } if (k == "" || v == nil) && k != "available_int" { continue } switch val := v.(type) { case string: // Sanitize: replace ' \n \r with space sanitized := sanitizeString(val) params[fmt.Sprintf("pp_%s", k)] = sanitized scripts = append(scripts, fmt.Sprintf("ctx._source.%s=params.pp_%s;", k, k)) case int, float64: scripts = append(scripts, fmt.Sprintf("ctx._source.%s=%v;", k, val)) case []interface{}: params[fmt.Sprintf("pp_%s", k)] = val scripts = append(scripts, fmt.Sprintf("ctx._source.%s=params.pp_%s;", k, k)) } } scriptSource := strings.Join(scripts, "") // Build update by query body updateBody := map[string]interface{}{ "query": boolQuery, "script": map[string]interface{}{ "source": scriptSource, "params": params, }, } bodyBytes, err := json.Marshal(updateBody) if err != nil { return fmt.Errorf("failed to marshal update body: %w", err) } // Execute update by query with refresh=true, slices=5, conflicts="proceed" refreshTrue := true req := esapi.UpdateByQueryRequest{ Index: []string{indexName}, Body: bytes.NewReader(bodyBytes), Refresh: &refreshTrue, Slices: 5, Conflicts: "proceed", } res, err := req.Do(ctx, e.client) if err != nil { common.Error("Failed to execute update by query", err) return fmt.Errorf("failed to execute update by query: %w", err) } defer res.Body.Close() if res.IsError() { bodyBytes, _ := io.ReadAll(res.Body) return fmt.Errorf("elasticsearch update by query error: %s, body: %s", res.Status(), string(bodyBytes)) } common.Debug("ElasticsearchConnection.updateChunksByQuery completed", zap.String("indexName", indexName)) return nil } // sanitizeString replaces ' \n \r with space func sanitizeString(s string) string { s = strings.ReplaceAll(s, "'", " ") s = strings.ReplaceAll(s, "\n", " ") s = strings.ReplaceAll(s, "\r", " ") return strings.TrimSpace(s) } // copyFields creates a shallow copy of a map func copyFields(m map[string]interface{}) map[string]interface{} { result := make(map[string]interface{}) for k, v := range m { result[k] = v } return result } // DeleteChunks deletes chunks from a dataset index by condition func (e *elasticsearchEngine) DeleteChunks(ctx context.Context, condition map[string]interface{}, indexName string, datasetID string) (int64, error) { // For ES, index name is just indexName (e.g., "ragflow_{tenantID}"), not indexName_datasetID fullIndexName := indexName common.Info("Deleting chunks from Elasticsearch index", zap.String("index_name", fullIndexName), zap.Any("condition", condition)) // Check if index exists exists, err := e.indexExists(ctx, fullIndexName) if err != nil { return 0, fmt.Errorf("failed to check index existence: %w", err) } if !exists { common.Warn(fmt.Sprintf("Index %s does not exist, skipping delete", fullIndexName)) return 0, nil } // Build bool query from condition var mustClauses []map[string]interface{} var filterClauses []map[string]interface{} var mustNotClauses []map[string]interface{} // Handle chunk IDs - use terms query on "id" field instead of ids query on _id if idVal, ok := condition["id"]; ok && idVal != nil { switch v := idVal.(type) { case []interface{}: ids := make([]string, 0, len(v)) for _, id := range v { if s, ok := id.(string); ok { ids = append(ids, s) } } if len(ids) > 0 { mustClauses = append(mustClauses, map[string]interface{}{ "terms": map[string]interface{}{"id": ids}, }) } case string: mustClauses = append(mustClauses, map[string]interface{}{ "term": map[string]interface{}{"id": v}, }) } } // Handle kb_id - add as term filter if kbID, ok := condition["kb_id"].(string); ok && kbID != "" { filterClauses = append(filterClauses, map[string]interface{}{ "term": map[string]interface{}{"kb_id": kbID}, }) } // Add all other conditions as filters/must/must_not for k, v := range condition { if k == "id" || k == "kb_id" { continue // Already handled above } if k == "exists" { filterClauses = append(filterClauses, map[string]interface{}{ "exists": map[string]interface{}{"field": v}, }) } else if k == "must_not" { if m, ok := v.(map[string]interface{}); ok { for kk, vv := range m { if kk == "exists" { mustNotClauses = append(mustNotClauses, map[string]interface{}{ "exists": map[string]interface{}{"field": vv}, }) } } } } else if v != nil { if listVal, ok := v.([]interface{}); ok { mustClauses = append(mustClauses, map[string]interface{}{ "terms": map[string]interface{}{k: listVal}, }) } else if _, ok := v.(string); ok { mustClauses = append(mustClauses, map[string]interface{}{ "term": map[string]interface{}{k: v}, }) } else if _, ok := v.(int); ok { mustClauses = append(mustClauses, map[string]interface{}{ "term": map[string]interface{}{k: v}, }) } } } // Build the query var qry map[string]interface{} if len(filterClauses) == 0 && len(mustClauses) == 0 && len(mustNotClauses) == 0 { qry = map[string]interface{}{"match_all": map[string]interface{}{}} } else { boolMap := map[string]interface{}{} if len(filterClauses) > 0 { boolMap["filter"] = filterClauses } if len(mustClauses) > 0 { boolMap["must"] = mustClauses } if len(mustNotClauses) > 0 { boolMap["must_not"] = mustNotClauses } qry = map[string]interface{}{"bool": boolMap} } // Build delete by query body deleteBody := map[string]interface{}{ "query": qry, } bodyBytes, err := json.Marshal(deleteBody) if err != nil { return 0, fmt.Errorf("failed to marshal delete body: %w", err) } // Execute delete by query with refresh=true refreshTrue := true req := esapi.DeleteByQueryRequest{ Index: []string{fullIndexName}, Body: bytes.NewReader(bodyBytes), Refresh: &refreshTrue, } res, err := req.Do(ctx, e.client) if err != nil { common.Error("Failed to execute delete by query", err) if strings.Contains(err.Error(), "not_found") { return 0, nil } return 0, fmt.Errorf("failed to execute delete by query: %w", err) } defer res.Body.Close() if res.IsError() { bodyBytes, _ := io.ReadAll(res.Body) errStr := string(bodyBytes) if strings.Contains(errStr, "not_found") { return 0, nil } common.Sugar.Errorw("Elasticsearch delete by query returned error", "status", res.Status()) return 0, fmt.Errorf("elasticsearch delete by query returned error: %s", res.Status()) } // Parse response var result map[string]interface{} if err := json.NewDecoder(res.Body).Decode(&result); err != nil { common.Error("Failed to parse delete response", err) return 0, fmt.Errorf("failed to parse delete response: %w", err) } deleted := int64(0) if d, ok := result["deleted"].(float64); ok { deleted = int64(d) } common.Info("Successfully deleted chunks", zap.String("index_name", fullIndexName), zap.Int64("deleted_count", deleted)) return deleted, nil } // SearchResponse Elasticsearch search response type SearchResponse struct { Hits struct { Total struct { Value int64 `json:"value"` } `json:"total"` Hits []struct { ID string `json:"_id"` Index string `json:"_index"` Score float64 `json:"_score"` Source map[string]interface{} `json:"_source"` Fields map[string]interface{} `json:"fields"` // ES 9.x stores dense_vector here Highlight map[string]interface{} `json:"highlight,omitempty"` // Sort is populated when the request body specifies a `sort` // clause. The last hit's Sort is the cursor for the next // search_after request — without it, deep pagination can't // advance. Sort []interface{} `json:"sort,omitempty"` } `json:"hits"` } `json:"hits"` Aggregations map[string]interface{} `json:"aggregations"` } // Search executes search with unified types.SearchRequest func (e *elasticsearchEngine) Search(ctx context.Context, req *types.SearchRequest) (*types.SearchResult, error) { types.LogSearchRequest("Elasticsearch", req) // Validate inputs and set defaults if len(req.IndexNames) == 0 { return nil, fmt.Errorf("index names cannot be empty") } offset := req.Offset if offset < 0 { offset = 0 } limit := req.Limit if limit <= 0 { limit = 30 } // Detect index types isSkillIndex := false for _, idx := range req.IndexNames { if strings.HasPrefix(idx, "skill_") { isSkillIndex = true break } } // Build bool query from condition boolQuery := buildBoolQueryFromCondition(req.Filter, req.KbIDs, isSkillIndex) // Extract vector_similarity_weight from FusionExpr var matchText *types.MatchTextExpr var matchDense *types.MatchDenseExpr vectorSimilarityWeight := 0.5 for _, expr := range req.MatchExprs { if expr == nil { continue } switch m := expr.(type) { case *types.FusionExpr: if m.Method == "weighted_sum" { if weights, ok := m.FusionParams["weights"].(string); ok { // Assert structure only when FusionExpr has weighted_sum with weights if len(req.MatchExprs) != 3 { return nil, fmt.Errorf("match_expressions must have exactly 3 elements with FusionExpr, got %d", len(req.MatchExprs)) } if _, ok := req.MatchExprs[0].(*types.MatchTextExpr); !ok { return nil, fmt.Errorf("match_expressions[0] must be MatchTextExpr") } if _, ok := req.MatchExprs[1].(*types.MatchDenseExpr); !ok { return nil, fmt.Errorf("match_expressions[1] must be MatchDenseExpr") } if _, ok := req.MatchExprs[2].(*types.FusionExpr); !ok { return nil, fmt.Errorf("match_expressions[2] must be FusionExpr") } parts := strings.Split(weights, ",") if len(parts) == 2 { if w, err := strconv.ParseFloat(parts[1], 64); err == nil { vectorSimilarityWeight = w } } } } case *types.MatchTextExpr: matchText = m case *types.MatchDenseExpr: matchDense = m } } // Build query body with text match and/or knn match queryBody := make(map[string]interface{}) if matchText != nil { textQuery := buildQueryStringQuery(matchText, vectorSimilarityWeight, isSkillIndex) if boolQuery != nil { if boolMap, ok := boolQuery["bool"].(map[string]interface{}); ok { if must, ok := boolMap["must"].([]interface{}); ok { must = append(must, textQuery) boolMap["must"] = must } else { boolMap["must"] = []interface{}{textQuery} } boolMap["boost"] = 1.0 - vectorSimilarityWeight } } else { boolQuery = textQuery } } hasVectorMatch := matchDense != nil && len(matchDense.EmbeddingData) > 0 if hasVectorMatch { k := matchDense.TopN if k <= 0 { k = limit } if k <= 0 { k = 1024 } numCandidates := k * 2 similarity := 0.0 if matchDense.ExtraOptions != nil { if sim, ok := matchDense.ExtraOptions["similarity"].(float64); ok { similarity = sim } } vectorFieldName := matchDense.VectorColumnName knnQuery := map[string]interface{}{ "field": vectorFieldName, "query_vector": matchDense.EmbeddingData, "k": k, "num_candidates": numCandidates, "similarity": similarity, "filter": boolQuery, } queryBody["knn"] = knnQuery if boolQuery != nil { queryBody["query"] = boolQuery } } else if boolQuery != nil { queryBody["query"] = boolQuery } else { queryBody["query"] = map[string]interface{}{ "match_all": map[string]interface{}{}, } } // Add rank_feature queries if req.RankFeature != nil && len(req.RankFeature) > 0 && !isSkillIndex { rankFeatureQuery := buildRankFeatureQuery(req.RankFeature) if rankFeatureQuery != nil { if boolQuery, ok := queryBody["query"].(map[string]interface{}); ok { if boolMap, ok := boolQuery["bool"].(map[string]interface{}); ok { if should, ok := boolMap["should"].([]interface{}); ok { for _, q := range rankFeatureQuery { boolMap["should"] = append(should, q) } } else { interfaceSlice := make([]interface{}, len(rankFeatureQuery)) for i, q := range rankFeatureQuery { interfaceSlice[i] = q } boolMap["should"] = interfaceSlice } } } } } // Add sorting if order_by specified if req.OrderBy != nil && len(req.OrderBy.Fields) > 0 { sort := parseOrderByExpr(req.OrderBy) if len(sort) > 0 { queryBody["sort"] = sort } } // Determine use_search_after for deep pagination // // ES rejects from + size combinations where from + size > index.max_result_window // (default 10,000) — see https://www.elastic.co/guide/en/elasticsearch/reference/current/paginate-search-results.html. // For those requests we must drop `from` and walk the result set with // `search_after` instead. The preconditions mirror the Python reference // (rag/utils/es_conn.py): // - explicit OrderBy is required (search_after needs a stable cursor, // and _score / KNN-similarity sorts are not unique enough to be safe) // - no dense/KNN match (knn queries do not honour `search_after` in the // same way and the Python path explicitly disallows them here) hasDense := hasVectorMatch hasExplicitSort := req.OrderBy != nil && len(req.OrderBy.Fields) > 0 useSearchAfter := limit > 0 && (offset+limit > common.MAX_RESULT_WINDOW) && hasExplicitSort && !hasDense // Apply offset/limit pagination. When useSearchAfter is true, the // caller is going to drive pagination via searchAfterCursor() // instead, so we must NOT emit from/size here — leaving them out // is the whole point of routing to the search_after path. if !useSearchAfter && limit > 0 { queryBody["size"] = limit queryBody["from"] = offset } // Set _source and fields for vector fields hasTextMatch := matchText != nil if len(req.SelectFields) > 0 { // Use caller-specified fields, add pagerank_fld/tag_fld if needed queryBody["_source"] = req.SelectFields if hasTextMatch || hasVectorMatch { if !isSkillIndex { if !slices.Contains(req.SelectFields, common.PAGERANK_FLD) { queryBody["_source"] = append(queryBody["_source"].([]string), common.PAGERANK_FLD) } if !slices.Contains(req.SelectFields, common.TAG_FLD) { queryBody["_source"] = append(queryBody["_source"].([]string), common.TAG_FLD) } } } var vectorFields []string for _, f := range req.SelectFields { if strings.HasSuffix(f, "_vec") { vectorFields = append(vectorFields, f) } } if len(vectorFields) > 0 { queryBody["fields"] = vectorFields } } else { // No explicit SelectFields - use match_all, but add pagerank_fld/tag_fld for scoring if needed if hasTextMatch || hasVectorMatch { if !isSkillIndex { queryBody["_source"] = []string{common.PAGERANK_FLD, common.TAG_FLD} } } } // Serialize query var buf bytes.Buffer if err := json.NewEncoder(&buf).Encode(queryBody); err != nil { return nil, fmt.Errorf("error encoding query: %w", err) } // Execute search. When useSearchAfter is true we must NOT send // from/size (we dropped them above) and instead walk the result set // page-by-page with the search_after cursor — ES otherwise returns // the first page and the caller gets the wrong page. var ( totalHits int64 allResults []map[string]interface{} err error ) if useSearchAfter { allResults, totalHits, err = e.searchAfterCursor(ctx, req, queryBody, offset, limit) if err != nil { return nil, err } } else { // WithBody takes an io.Reader that the Go client streams // directly into the request. Reusing &buf across iterations // would drain it on the first request and leave the rest // with an empty body — so we copy the bytes once and hand // each iteration a fresh bytes.NewReader. payload := append([]byte(nil), buf.Bytes()...) for _, indexName := range req.IndexNames { res, err := e.client.Search( e.client.Search.WithContext(ctx), e.client.Search.WithIndex(indexName), e.client.Search.WithBody(bytes.NewReader(payload)), e.client.Search.WithTrackTotalHits(true), ) if err != nil { common.Warn("Elasticsearch query failed", zap.String("index", indexName), zap.Error(err)) continue } defer res.Body.Close() if res.IsError() { bodyBytes, _ := io.ReadAll(res.Body) common.Warn("Elasticsearch error response", zap.String("index", indexName), zap.String("body", string(bodyBytes))) continue } // Parse response and return results var esResp SearchResponse if err := json.NewDecoder(res.Body).Decode(&esResp); err != nil { common.Warn("Elasticsearch failed to parse response", zap.String("index", indexName), zap.Error(err)) continue } searchChunks := convertESResponse(&esResp, "") totalHits += esResp.Hits.Total.Value allResults = append(allResults, searchChunks...) } } // Post-processing: Sort results by score if len(allResults) > 0 && (matchText != nil || hasVectorMatch) { scoreColumn := "_score" if matchText != nil && hasVectorMatch { scoreColumn = "SCORE" } pagerankField := common.PAGERANK_FLD if isSkillIndex { pagerankField = "" } allResults = calculateScores(allResults, scoreColumn, pagerankField) allResults = sortByScore(allResults, limit) } common.Info("ES Search completed", zap.Int("returnedRows", len(allResults)), zap.Int64("totalHits", totalHits)) return &types.SearchResult{ Chunks: allResults, Total: totalHits, }, nil } // searchAfterFetcher issues one ES search request with the given batch // size and search_after cursor, returning the decoded response. Defined // as a function type so the pagination logic below can be unit-tested // with a mock fetcher instead of a real Elasticsearch client. type searchAfterFetcher func( ctx context.Context, baseQuery map[string]interface{}, batch int, cursor []interface{}, trackTotalHits bool, ) (SearchResponse, error) // searchAfterCursor walks ES with the search_after pagination protocol, // returning the page [offset, offset+limit) of an explicitly-sorted // result set. Used when offset+limit exceeds common.MAX_RESULT_WINDOW // and ES would otherwise reject the from/size combination. // // Mirrors rag/utils/es_conn.py:ESConnection._search_with_search_after: // // 1. Drop from/size from the base query (the caller has already omitted // them on this path; this is a defensive no-op). // 2. Skip phase: discard hits until we have skipped `offset` of them. // 3. Take phase: collect hits until we have `limit` of them, or the // index is exhausted. // 4. After each batch, advance the cursor with the last hit's `sort` // field. If `sort` is missing or unchanged, the index is exhausted. // // The first request carries trackTotalHits=true so the caller still // gets an accurate total; subsequent requests skip it for efficiency. // Returns the (possibly empty) collected hits and the total hit count // from the first response. func (e *elasticsearchEngine) searchAfterCursor( ctx context.Context, req *types.SearchRequest, baseQuery map[string]interface{}, offset, limit int, ) ([]map[string]interface{}, int64, error) { // Defensive: strip from/size if the caller left them in. In the // current code path they are never set when useSearchAfter is true, // but the base query is a shared map and future callers may forget. delete(baseQuery, "from") delete(baseQuery, "size") return searchAfterPaginate(ctx, baseQuery, offset, limit, e.buildSearchAfterFetcher(req)) } // buildSearchAfterFetcher returns a fetcher that delegates each // iteration to executeSearchRequest, which talks to the real ES client. func (e *elasticsearchEngine) buildSearchAfterFetcher(req *types.SearchRequest) searchAfterFetcher { return func( ctx context.Context, baseQuery map[string]interface{}, batch int, cursor []interface{}, trackTotalHits bool, ) (SearchResponse, error) { return e.executeSearchRequest(ctx, req, baseQuery, batch, cursor, trackTotalHits) } } // searchAfterPaginate is the pure, callback-driven pagination loop // shared by the engine and the unit tests. See searchAfterCursor for // the semantics. func searchAfterPaginate( ctx context.Context, baseQuery map[string]interface{}, offset, limit int, fetch searchAfterFetcher, ) ([]map[string]interface{}, int64, error) { var ( cursor []interface{} totalHits int64 collected []map[string]interface{} collectedTake int firstCall = true ) // Skip phase: walk past `offset` hits without retaining them. remainingSkip := offset for remainingSkip > 0 { batch := remainingSkip if batch > common.SearchAfterBatchSize { batch = common.SearchAfterBatchSize } resp, err := fetch(ctx, baseQuery, batch, cursor, firstCall) firstCall = false if err != nil { return nil, 0, err } if totalHits == 0 { totalHits = resp.Hits.Total.Value } if len(resp.Hits.Hits) == 0 { break } nextCursor := resp.Hits.Hits[len(resp.Hits.Hits)-1].Sort if len(nextCursor) == 0 || sortValuesEqual(nextCursor, cursor) { // ES returned hits but no usable cursor (e.g. sort field // missing or unchanged). The index is exhausted from our // point of view. break } cursor = nextCursor remainingSkip -= len(resp.Hits.Hits) if len(resp.Hits.Hits) < batch { // Short batch — we asked for more than was available, so // the cursor is at the end of the index. break } } // Take phase: collect up to `limit` hits. ES may return up to // `batch` hits per request, but we stop at `limit` (the absolute // target) regardless of how many we asked for in this iteration. for collectedTake < limit { want := limit - collectedTake batch := want if batch > common.SearchAfterBatchSize { batch = common.SearchAfterBatchSize } resp, err := fetch(ctx, baseQuery, batch, cursor, firstCall) firstCall = false if err != nil { return nil, 0, err } if totalHits == 0 { totalHits = resp.Hits.Total.Value } if len(resp.Hits.Hits) == 0 { break } // Convert and append. We could parallelize the conversion with // the next request, but conversion is cheap relative to the // ES round-trip, so keep the loop straightforward. for _, hit := range resp.Hits.Hits { if collectedTake >= limit { break } chunk := hit.Source if chunk == nil { chunk = map[string]interface{}{} } chunk["_score"] = hit.Score chunk["_id"] = hit.ID chunk["_index"] = hit.Index collected = append(collected, chunk) collectedTake++ } // Reached the absolute limit — stop without advancing the // cursor (we already have what was asked for). if collectedTake >= limit { break } nextCursor := resp.Hits.Hits[len(resp.Hits.Hits)-1].Sort if len(nextCursor) == 0 || sortValuesEqual(nextCursor, cursor) { break } cursor = nextCursor if len(resp.Hits.Hits) < batch { break } } // If we never sent a request (e.g. offset == 0 and limit == 0) we // still need a total. Issue one count-only request. if totalHits == 0 { resp, err := fetch(ctx, baseQuery, 0, nil, true) if err != nil { return nil, 0, err } totalHits = resp.Hits.Total.Value } return collected, totalHits, nil } // executeSearchRequest sends one ES search request with the given // batch size and search_after cursor. If trackTotalHits is true the // request asks ES to compute an exact total (cheap to omit on // pagination iterations after the first). func (e *elasticsearchEngine) executeSearchRequest( ctx context.Context, req *types.SearchRequest, baseQuery map[string]interface{}, batch int, cursor []interface{}, trackTotalHits bool, ) (SearchResponse, error) { queryBody := make(map[string]interface{}, len(baseQuery)+2) for k, v := range baseQuery { queryBody[k] = v } if batch > 0 { queryBody["size"] = batch } if len(cursor) > 0 { queryBody["search_after"] = cursor } var buf bytes.Buffer if err := json.NewEncoder(&buf).Encode(queryBody); err != nil { return SearchResponse{}, fmt.Errorf("error encoding query: %w", err) } res, err := e.client.Search( e.client.Search.WithContext(ctx), e.client.Search.WithIndex(req.IndexNames...), e.client.Search.WithBody(&buf), e.client.Search.WithTrackTotalHits(trackTotalHits), ) if err != nil { return SearchResponse{}, fmt.Errorf("elasticsearch search failed: %w", err) } defer res.Body.Close() if res.IsError() { bodyBytes, _ := io.ReadAll(res.Body) return SearchResponse{}, fmt.Errorf("elasticsearch error response: %s", string(bodyBytes)) } var esResp SearchResponse if err := json.NewDecoder(res.Body).Decode(&esResp); err != nil { return SearchResponse{}, fmt.Errorf("elasticsearch failed to parse response: %w", err) } return esResp, nil } // sortValuesEqual reports whether two sort cursors are identical. // ES guarantees that successive requests with `search_after: ` // advance strictly past the cursor, so an unchanged cursor between // iterations means the index is exhausted. func sortValuesEqual(a, b []interface{}) bool { if len(a) != len(b) { return false } for i := range a { if a[i] != b[i] { return false } } return true } // buildBoolQueryFromCondition builds an ES bool query from condition map // For skill index, uses 'status' field instead of 'available_int' func buildBoolQueryFromCondition(filter map[string]interface{}, kbIDs []string, isSkillIndex bool) map[string]interface{} { var mustClauses []interface{} var filterClauses []interface{} var shouldClauses []interface{} // Add kb_id to condition if kbIDs != nil && len(kbIDs) > 0 { filterClauses = append(filterClauses, map[string]interface{}{ "terms": map[string]interface{}{"kb_id": kbIDs}, }) } // For skill index, add status = "1" filter by default (active skills) if isSkillIndex { filterClauses = append(filterClauses, map[string]interface{}{ "term": map[string]interface{}{ "status": "1", }, }) } if filter == nil { filter = make(map[string]interface{}) } for k, v := range filter { // For skill index, handle 'status' field instead of 'available_int' if isSkillIndex && k == "status" { if v == nil || v == "" { continue } if listVal, ok := v.([]interface{}); ok && len(listVal) > 0 { filterClauses = append(filterClauses, map[string]interface{}{ "terms": map[string]interface{}{"status": listVal}, }) } else if strVal, ok := v.(string); ok && strVal != "" { filterClauses = append(filterClauses, map[string]interface{}{ "term": map[string]interface{}{"status": strVal}, }) } continue } if k == "available_int" { var numVal float64 switch val := v.(type) { case float64: numVal = val case int: numVal = float64(val) case int64: numVal = float64(val) default: continue } if numVal == 0 { filterClauses = append(filterClauses, map[string]interface{}{ "range": map[string]interface{}{"available_int": map[string]interface{}{"lt": 1}}, }) } else { filterClauses = append(filterClauses, map[string]interface{}{ "bool": map[string]interface{}{ "must_not": []map[string]interface{}{ {"range": map[string]interface{}{"available_int": map[string]interface{}{"lt": 1}}}, }, }, }) } continue } if k == "id" { if v == nil || v == "" { continue } if listVal, ok := v.([]interface{}); ok && len(listVal) > 0 { shouldClauses = append(shouldClauses, map[string]interface{}{"terms": map[string]interface{}{"id": listVal}}, map[string]interface{}{"terms": map[string]interface{}{"_id": listVal}}, ) } else if strVal, ok := v.(string); ok && strVal != "" { shouldClauses = append(shouldClauses, map[string]interface{}{"term": map[string]interface{}{"id": strVal}}, map[string]interface{}{"term": map[string]interface{}{"_id": strVal}}, ) } else if intVal, ok := v.(int); ok && intVal != 0 { shouldClauses = append(shouldClauses, map[string]interface{}{"term": map[string]interface{}{"id": intVal}}, map[string]interface{}{"term": map[string]interface{}{"_id": intVal}}, ) } continue } if v == nil || v == "" { continue } if listVal, ok := v.([]interface{}); ok { filterClauses = append(filterClauses, map[string]interface{}{ "terms": map[string]interface{}{k: listVal}, }) } else if strListVal, ok := v.([]string); ok { filterClauses = append(filterClauses, map[string]interface{}{ "terms": map[string]interface{}{k: strListVal}, }) } else if strVal, ok := v.(string); ok && strVal != "" { filterClauses = append(filterClauses, map[string]interface{}{ "term": map[string]interface{}{k: strVal}, }) } else if intVal, ok := v.(int); ok { filterClauses = append(filterClauses, map[string]interface{}{ "term": map[string]interface{}{k: intVal}, }) } else if floatVal, ok := v.(float64); ok { filterClauses = append(filterClauses, map[string]interface{}{ "term": map[string]interface{}{k: floatVal}, }) } } // Build the bool query boolQuery := make(map[string]interface{}) if len(mustClauses) > 0 { boolQuery["must"] = mustClauses } if len(filterClauses) > 0 { boolQuery["filter"] = filterClauses } if len(shouldClauses) > 0 { boolQuery["should"] = shouldClauses boolQuery["minimum_should_match"] = 1 } if len(boolQuery) == 0 { return nil } return map[string]interface{}{"bool": boolQuery} } // buildQueryStringQuery builds a query_string query from MatchTextExpr // When isSkillIndex is true, uses skill-specific fields (name_tks, tags_tks, etc.) // Otherwise uses document fields (title_tks, content_ltks, etc.) func buildQueryStringQuery(matchText *types.MatchTextExpr, vectorSimilarityWeight float64, isSkillIndex bool) map[string]interface{} { if matchText == nil { return nil } minimumShouldMatch := "0%" if matchText.ExtraOptions != nil { if msm, ok := matchText.ExtraOptions["minimum_should_match"].(float64); ok { minimumShouldMatch = fmt.Sprintf("%d%%", int(msm*100)) } } fields := matchText.Fields if fields == nil || len(fields) == 0 { if isSkillIndex { fields = []string{"name_tks^10", "tags_tks^5", "description_tks^3", "content_tks^1"} } else { fields = []string{"title_tks^10", "title_sm_tks^5", "important_kwd^30", "important_tks^20", "question_tks^20", "content_ltks^2", "content_sm_ltks"} } } boost := 1.0 if matchText.ExtraOptions != nil { if b, ok := matchText.ExtraOptions["boost"].(float64); ok { boost = b } } return map[string]interface{}{ "query_string": map[string]interface{}{ "fields": fields, "type": "best_fields", "query": matchText.MatchingText, "minimum_should_match": minimumShouldMatch, "boost": boost, }, } } // buildRankFeatureQuery builds rank_feature queries for learning to rank func buildRankFeatureQuery(rankFeature map[string]float64) []map[string]interface{} { if rankFeature == nil || len(rankFeature) == 0 { return nil } // Sort keys for deterministic query order (Go map iteration is randomized) keys := make([]string, 0, len(rankFeature)) for k := range rankFeature { keys = append(keys, k) } sort.Strings(keys) var queries []map[string]interface{} for _, fld := range keys { if fld == common.PAGERANK_FLD { continue } sc := rankFeature[fld] tagField := fmt.Sprintf("%s.%s", common.TAG_FLD, fld) queries = append(queries, map[string]interface{}{ "rank_feature": map[string]interface{}{ "field": tagField, "linear": map[string]interface{}{}, "boost": sc, }, }) } return queries } // GetChunk gets a chunk by ID using ES search API func (e *elasticsearchEngine) GetChunk(ctx context.Context, baseName, chunkID string, datasetIDs []string) (interface{}, error) { // Try search by doc_id field (which is stored in the document) for _, datasetID := range datasetIDs { searchReq := map[string]interface{}{ "query": map[string]interface{}{ "bool": map[string]interface{}{ "must": []map[string]interface{}{ {"term": map[string]interface{}{"id": chunkID}}, {"term": map[string]interface{}{"kb_id": datasetID}}, }, }, }, } body, err := json.Marshal(searchReq) if err != nil { return nil, fmt.Errorf("failed to marshal search request: %w", err) } res, err := e.client.Search( e.client.Search.WithContext(ctx), e.client.Search.WithIndex(baseName), e.client.Search.WithBody(bytes.NewReader(body)), ) if err != nil { return nil, fmt.Errorf("failed to search for chunk: %w", err) } if res.IsError() { res.Body.Close() return nil, fmt.Errorf("failed to search for chunk: %s", res.Status()) } var searchResult map[string]interface{} if err := json.NewDecoder(res.Body).Decode(&searchResult); err != nil { res.Body.Close() return nil, fmt.Errorf("failed to parse search response: %w", err) } res.Body.Close() hits, ok := searchResult["hits"].(map[string]interface{}) if !ok { continue } hitList, ok := hits["hits"].([]interface{}) if !ok || len(hitList) == 0 { continue } firstHit, ok := hitList[0].(map[string]interface{}) if !ok { continue } source, ok := firstHit["_source"].(map[string]interface{}) if !ok { continue } common.Info("GetChunk found hit", zap.String("baseName", baseName), zap.String("chunkID", chunkID)) source["id"] = chunkID return source, nil } common.Info("GetChunk no hits found", zap.String("baseName", baseName), zap.String("chunkID", chunkID)) return nil, nil } // GetFields extracts the requested fields from ES search response chunks // // Unlike Infinity, Elasticsearch does NOT use convertSelectFields before querying. // The original requested field names ARE the database column names: // - "content_with_weight" is stored and returned as "content_with_weight" // - No field name mapping is needed in GetFields func (e *elasticsearchEngine) GetFields(chunks []map[string]interface{}, fields []string) map[string]map[string]interface{} { common.Info("GetFields called", zap.Int("chunkCount", len(chunks)), zap.Strings("fields", fields)) result := make(map[string]map[string]interface{}) if len(fields) == 0 || len(chunks) == 0 { return result } // Build field set for lookup fieldSet := make(map[string]bool) for _, f := range fields { fieldSet[f] = true } for _, chunk := range chunks { docID, ok := elasticsearchChunkID(chunk) if !ok { continue } if id, ok := chunk["id"].(string); !ok || id == "" { chunk["id"] = docID } m := make(map[string]interface{}) for field := range fieldSet { val := chunk[field] if val == nil { continue } if listVal, ok := val.([]interface{}); ok { if len(listVal) == 1 { if _, isArray := listVal[0].([]interface{}); !isArray { val = listVal[0] } } } if _, ok := val.([]interface{}); ok { m[field] = val continue } if field == "available_int" { if _, ok := val.(int); ok { m[field] = val continue } if _, ok := val.(float64); ok { m[field] = val continue } } if _, ok := val.(string); !ok { val = fmt.Sprintf("%v", val) } m[field] = val } if len(m) > 0 { result[docID] = m } } common.Info("GetFields result", zap.Int("resultCount", len(result)), zap.Strings("keys", func() []string { keys := make([]string, 0, len(result)) for k := range result { keys = append(keys, k) } return keys }())) return result } // GetAggregation aggregates chunk values by field name // Input: [{"docnm_kwd": "docA"}, {"docnm_kwd": "docA"}, {"docnm_kwd": "docB"}] // Returns: [{"key": "docA", "count": 2}, {"key": "docB", "count": 1}] func (e *elasticsearchEngine) GetAggregation(chunks []map[string]interface{}, fieldName string) []map[string]interface{} { if len(chunks) == 0 || fieldName == "" { return []map[string]interface{}{} } tagCounts := make(map[string]int) for _, chunk := range chunks { value, ok := chunk[fieldName] if !ok || value == nil { continue } if valueStr, ok := value.(string); ok { if valueStr == "" { continue } separator := "," if fieldName == "tag_kwd" && strings.Contains(valueStr, "###") { separator = "###" } for _, tag := range strings.Split(valueStr, separator) { countElasticsearchAggregationTag(tagCounts, tag) } continue } if valueList, ok := value.([]interface{}); ok { for _, item := range valueList { if itemStr, ok := item.(string); ok { countElasticsearchAggregationTag(tagCounts, itemStr) } } } } if len(tagCounts) == 0 { return []map[string]interface{}{} } tags := make([]string, 0, len(tagCounts)) for tag := range tagCounts { tags = append(tags, tag) } slices.SortFunc(tags, func(a, b string) int { if byCount := cmp.Compare(tagCounts[b], tagCounts[a]); byCount != 0 { return byCount } return cmp.Compare(a, b) }) result := make([]map[string]interface{}, len(tags)) for i, tag := range tags { result[i] = map[string]interface{}{"key": tag, "count": tagCounts[tag]} } return result } // GetChunkIDs extracts chunk IDs from ES search response chunks. // Uses _id field (composite: {doc_id}_{kb_id}_{chunk_id}). func (e *elasticsearchEngine) GetChunkIDs(chunks []map[string]interface{}) []string { ids := make([]string, 0, len(chunks)) for _, chunk := range chunks { if id, ok := elasticsearchChunkID(chunk); ok { ids = append(ids, id) } } return ids } // GetHighlight returns highlighted text for matching keywords func (e *elasticsearchEngine) GetHighlight(chunks []map[string]interface{}, keywords []string, fieldName string) map[string]string { result := make(map[string]string) if len(chunks) == 0 || len(keywords) == 0 { return result } normalizedKeywords := normalizeElasticsearchHighlightKeywords(keywords) englishPatterns := compileElasticsearchHighlightPatterns(normalizedKeywords) nonEnglishPattern := compileElasticsearchNonEnglishHighlightPattern(normalizedKeywords) for _, chunk := range chunks { docID, ok := elasticsearchChunkID(chunk) if !ok { continue } if highlightText := firstElasticsearchHighlight(chunk); highlightText != "" { result[docID] = highlightText continue } txt, ok := chunk[fieldName].(string) if fieldName == "content_with_weight" && (!ok || txt == "") { txt, ok = chunk["content"].(string) } if !ok || txt == "" { continue } if elasticsearchHighlightEmTagRE.MatchString(txt) { result[docID] = txt continue } txt = elasticsearchHighlightNewlineRE.ReplaceAllString(txt, " ") segments := elasticsearchHighlightDelimiterRE.Split(txt, -1) var highlightedSegments []string for _, segment := range segments { segmentToCheck := segment if isMostlyEnglishElasticsearchSegment(segment) { for _, pattern := range englishPatterns { segmentToCheck = pattern.ReplaceAllString(segmentToCheck, "$1$2$3") } } else if nonEnglishPattern != nil { segmentToCheck = nonEnglishPattern.ReplaceAllStringFunc(segmentToCheck, func(match string) string { return "" + match + "" }) } if segmentToCheck != segment { highlightedSegments = append(highlightedSegments, strings.TrimSpace(segmentToCheck)) } } if len(highlightedSegments) > 0 { result[docID] = strings.Join(highlightedSegments, "... ") } } return result } func elasticsearchChunkID(chunk map[string]interface{}) (string, bool) { if id, ok := chunk["id"].(string); ok && id != "" { return id, true } if id, ok := chunk["_id"].(string); ok && id != "" { return id, true } return "", false } func firstElasticsearchHighlight(chunk map[string]interface{}) string { highlight, ok := chunk["highlight"].(map[string]interface{}) if !ok || len(highlight) == 0 { return "" } for _, vals := range highlight { if arr, ok := vals.([]interface{}); ok && len(arr) > 0 { if str, ok := arr[0].(string); ok { return str } } } return "" } func countElasticsearchAggregationTag(counts map[string]int, tag string) { if tag = strings.TrimSpace(tag); tag != "" { counts[tag]++ } } func isMostlyEnglishElasticsearchSegment(segment string) bool { totalCount := len(elasticsearchLetterRE.FindAllString(segment, -1)) return totalCount > 0 && float64(len(elasticsearchEnglishLetterRE.FindAllString(segment, -1)))/float64(totalCount) > 0.5 } func compileElasticsearchHighlightPatterns(keywords []string) []*regexp.Regexp { patterns := make([]*regexp.Regexp, 0, len(keywords)) for _, kw := range keywords { patterns = append(patterns, regexp.MustCompile(`(?i)(^|[ .?/'\"\(\)!,:;-])(`+regexp.QuoteMeta(kw)+`)([ .?/'\"\(\)!,:;-]|$)`)) } return patterns } func compileElasticsearchNonEnglishHighlightPattern(keywords []string) *regexp.Regexp { if len(keywords) == 0 { return nil } parts := make([]string, 0, len(keywords)) for _, kw := range keywords { parts = append(parts, regexp.QuoteMeta(kw)) } return regexp.MustCompile(strings.Join(parts, "|")) } func normalizeElasticsearchHighlightKeywords(keywords []string) []string { seen := make(map[string]struct{}, len(keywords)) normalized := make([]string, 0, len(keywords)) for _, kw := range keywords { if kw == "" { continue } if _, ok := seen[kw]; !ok { seen[kw] = struct{}{} normalized = append(normalized, kw) } } slices.SortStableFunc(normalized, func(a, b string) int { return cmp.Compare(len(b), len(a)) }) return normalized } // DropChunkStore deletes a chunk index func (e *elasticsearchEngine) DropChunkStore(ctx context.Context, baseName, datasetID string) error { return e.dropIndex(ctx, baseName) } // ChunkStoreExists checks if a chunk index exists func (e *elasticsearchEngine) ChunkStoreExists(ctx context.Context, baseName, datasetID string) (bool, error) { return e.indexExists(ctx, baseName) } // KNNScores performs a second-pass KNN search to get clean cosine similarities for ES. // This keeps chunk vectors in the index and asks ES to compute the cosine similarity. func (e *elasticsearchEngine) KNNScores(ctx context.Context, chunks []map[string]interface{}, queryVector []float64, topK int) (map[string]interface{}, error) { if len(chunks) == 0 || len(queryVector) == 0 { return nil, nil } // Extract chunk IDs from first search results chunkIDs := make([]string, 0, len(chunks)) for _, chunk := range chunks { if id, ok := chunk["_id"].(string); ok { chunkIDs = append(chunkIDs, id) } } if len(chunkIDs) == 0 { return nil, nil } common.Info("KNNScores starting", zap.Int("chunkCount", len(chunkIDs)), zap.Strings("chunkIDs", chunkIDs), zap.Int("vectorSize", len(queryVector))) // Build KNN-only query filtered by chunk IDs vectorSize := len(queryVector) k := len(chunkIDs) knnQuery := map[string]interface{}{ "field": fmt.Sprintf("q_%d_vec", vectorSize), "query_vector": queryVector, "k": k, "num_candidates": k * 2, "similarity": 0.0, // No threshold - get all "filter": map[string]interface{}{ "terms": map[string]interface{}{"id": chunkIDs}, }, } queryBody := map[string]interface{}{ "knn": knnQuery, "size": k, "_source": false, // Don't need source fields, only need _id and _score } body, err := json.Marshal(queryBody) if err != nil { return nil, fmt.Errorf("failed to marshal KNN query: %w", err) } //common.Info("KNNScores query body", zap.String("body", string(body))) // Execute search - use first index name from chunks if available indexName := "" if len(chunks) > 0 { if idx, ok := chunks[0]["_index"].(string); ok { indexName = idx } } res, err := e.client.Search( e.client.Search.WithContext(ctx), e.client.Search.WithIndex(indexName), e.client.Search.WithBody(bytes.NewReader(body)), ) if err != nil { return nil, fmt.Errorf("KNN scores search failed: %w", err) } defer res.Body.Close() if res.IsError() { bodyBytes, _ := io.ReadAll(res.Body) return nil, fmt.Errorf("KNN scores search returned error: %s, body: %s", res.Status(), string(bodyBytes)) } var esResp SearchResponse if err := json.NewDecoder(res.Body).Decode(&esResp); err != nil { return nil, fmt.Errorf("failed to parse KNN scores response: %w", err) } common.Info("KNNScores ES response", zap.Int("hitCount", len(esResp.Hits.Hits)), zap.Any("firstHit", func() interface{} { if len(esResp.Hits.Hits) > 0 { return esResp.Hits.Hits[0] } return nil }())) // Return raw ES response // Caller will pass to GetScores to extract scores knnResult := make(map[string]interface{}) knnResult["hits"] = map[string]interface{}{ "hits": esResp.Hits.Hits, } return knnResult, nil } // GetScores extracts similarity scores from KNN search result func (e *elasticsearchEngine) GetScores(knnResult map[string]interface{}) map[string]float64 { scores := make(map[string]float64) hits, ok := knnResult["hits"].(map[string]interface{}) if !ok { return scores } hitsList, ok := hits["hits"] if !ok { return scores } switch v := hitsList.(type) { case []interface{}: for _, h := range v { if hit, ok := h.(map[string]interface{}); ok { if docID, ok := hit["_id"].(string); ok && docID != "" { if scoreVal := hit["_score"]; scoreVal != nil { if score, ok := scoreVal.(float64); ok { scores[docID] = score } } } } } case []map[string]interface{}: for _, hit := range v { if docID, ok := hit["_id"].(string); ok && docID != "" { if scoreVal := hit["_score"]; scoreVal != nil { if score, ok := scoreVal.(float64); ok { scores[docID] = score } } } } default: // Handle slice of structs via reflection rv := reflect.ValueOf(v) if rv.Kind() == reflect.Slice { for i := 0; i < rv.Len(); i++ { elem := rv.Index(i) idField := elem.FieldByName("ID") if !idField.IsValid() { idField = elem.FieldByName("Id") } if !idField.IsValid() || idField.Kind() != reflect.String { continue } docID := idField.String() if docID == "" { continue } scoreField := elem.FieldByName("Score") if !scoreField.IsValid() || scoreField.Kind() != reflect.Float64 { continue } scores[docID] = scoreField.Float() } } } return scores } // loadSkillMapping loads the skill index mapping from config file func loadSkillMapping() (map[string]interface{}, error) { // Try multiple possible locations for the mapping file possiblePaths := []string{ "conf/skill_es_mapping.json", "../conf/skill_es_mapping.json", "/app/conf/skill_es_mapping.json", } var data []byte var err error for _, path := range possiblePaths { data, err = os.ReadFile(path) if err == nil { break } } if err != nil { // Fallback to default skill mapping if file not found return getDefaultSkillMapping(), nil } var mapping map[string]interface{} if err := json.Unmarshal(data, &mapping); err != nil { return nil, fmt.Errorf("failed to parse skill mapping: %w", err) } return mapping, nil } // getDefaultSkillMapping returns the default skill index mapping func getDefaultSkillMapping() map[string]interface{} { return map[string]interface{}{ "settings": map[string]interface{}{ "index": map[string]interface{}{ "number_of_shards": 1, "number_of_replicas": 0, "refresh_interval": "1000ms", }, }, "mappings": map[string]interface{}{ "dynamic": false, "properties": map[string]interface{}{ "skill_id": map[string]interface{}{ "type": "keyword", "store": true, }, "name": map[string]interface{}{ "type": "text", "index": false, "store": true, }, "name_tks": map[string]interface{}{ "type": "text", "analyzer": "whitespace", "store": true, }, "tags": map[string]interface{}{ "type": "text", "index": false, "store": true, }, "tags_tks": map[string]interface{}{ "type": "text", "analyzer": "whitespace", "store": true, }, "description": map[string]interface{}{ "type": "text", "index": false, "store": true, }, "description_tks": map[string]interface{}{ "type": "text", "analyzer": "whitespace", "store": true, }, "content": map[string]interface{}{ "type": "text", "index": false, "store": true, }, "content_tks": map[string]interface{}{ "type": "text", "analyzer": "whitespace", "store": true, }, "q_3072_vec": map[string]interface{}{ "type": "dense_vector", "dims": 3072, "index": true, "similarity": "cosine", }, "q_2560_vec": map[string]interface{}{ "type": "dense_vector", "dims": 2560, "index": true, "similarity": "cosine", }, "q_1536_vec": map[string]interface{}{ "type": "dense_vector", "dims": 1536, "index": true, "similarity": "cosine", }, "q_1024_vec": map[string]interface{}{ "type": "dense_vector", "dims": 1024, "index": true, "similarity": "cosine", }, "q_768_vec": map[string]interface{}{ "type": "dense_vector", "dims": 768, "index": true, "similarity": "cosine", }, "q_512_vec": map[string]interface{}{ "type": "dense_vector", "dims": 512, "index": true, "similarity": "cosine", }, "q_256_vec": map[string]interface{}{ "type": "dense_vector", "dims": 256, "index": true, "similarity": "cosine", }, "version": map[string]interface{}{ "type": "keyword", "store": true, }, "status": map[string]interface{}{ "type": "keyword", "store": true, }, "create_time": map[string]interface{}{ "type": "long", "store": true, }, "update_time": map[string]interface{}{ "type": "long", "store": true, }, }, }, } } // rerankWindow returns the candidate-window size shared by retrieval's // block fetch and slice. Mirrors Dealer._rerank_window in rag/nlp/search.py. // // `size` is the per-page size; the window MUST be an exact multiple of it, // otherwise the block fetched (offset // window) and the in-block page slice // (offset % window) drift apart and deep pagination silently drops results. // // The window targets a provider-friendly pool of ~64 candidates, bounded by // `topK` when given (i.e. when an external reranker is active), and is always // rounded UP to a whole number of pages to preserve the alignment invariant. func rerankWindow(size, topK int) int { if size <= 1 { if topK > 0 { return min(30, topK) } return 30 } window := ((64 + size - 1) / size) * size // ceil(64/size) * size if topK > 0 { if aligned := ((topK + size - 1) / size) * size; window > aligned { window = aligned } } return window } // calculatePagination calculates offset and limit based on page, size and topK func calculatePagination(page, size, topK int) (int, int) { if page < 1 { page = 1 } if size <= 0 { size = 30 } if topK <= 0 { topK = 1024 } window := rerankWindow(size, topK) offset := (page - 1) * window if offset < 0 { offset = 0 } return offset, window } // convertESResponse converts ES SearchResponse to unified chunks format func convertESResponse(esResp *SearchResponse, vectorFieldName string) []map[string]interface{} { if esResp == nil || esResp.Hits.Hits == nil { return []map[string]interface{}{} } chunks := make([]map[string]interface{}, len(esResp.Hits.Hits)) for i, hit := range esResp.Hits.Hits { chunks[i] = hit.Source if chunks[i] == nil { chunks[i] = make(map[string]interface{}) } chunks[i]["_score"] = hit.Score chunks[i]["_id"] = hit.ID chunks[i]["_index"] = hit.Index if len(hit.Highlight) > 0 { chunks[i]["highlight"] = hit.Highlight } } return chunks } // parseOrderByExpr parses the OrderBy expression into ES sort format func parseOrderByExpr(orderBy *types.OrderByExpr) []map[string]interface{} { if orderBy == nil || len(orderBy.Fields) == 0 { return nil } var result []map[string]interface{} for _, field := range orderBy.Fields { direction := "asc" if field.Type == types.SortDesc { direction = "desc" } // Skip id field (cannot order by text field) if field.Field == "id" { continue } // Special handling for page_num_int and top_int if field.Field == "page_num_int" || field.Field == "top_int" { result = append(result, map[string]interface{}{ field.Field: map[string]interface{}{ "order": direction, "unmapped_type": "float", "mode": "avg", "numeric_type": "double", }, }) } else if strings.HasSuffix(field.Field, "_int") || strings.HasSuffix(field.Field, "_flt") { // Fields ending with _int or _flt result = append(result, map[string]interface{}{ field.Field: map[string]interface{}{ "order": direction, "unmapped_type": "float", }, }) } else if field.Field == "_score" || field.Field == "score" { result = append(result, map[string]interface{}{ "_score": direction, }) } else { // Default: unmapped_type = keyword result = append(result, map[string]interface{}{ field.Field: map[string]interface{}{ "order": direction, "unmapped_type": "keyword", }, }) } } return result } // calculateScores calculates _score for chunks func calculateScores(chunks []map[string]interface{}, scoreColumn, pagerankField string) []map[string]interface{} { for i := range chunks { score := 0.0 if scoreVal, ok := chunks[i][scoreColumn]; ok { if f, ok := toFloat64(scoreVal); ok { score += f } } if pagerankField != "" { if prVal, ok := chunks[i][pagerankField]; ok { if f, ok := toFloat64(prVal); ok { score += f } } } chunks[i]["_score"] = score } return chunks } // toFloat64 converts a value to float64 func toFloat64(v interface{}) (float64, bool) { switch val := v.(type) { case float64: return val, true case float32: return float64(val), true case int: return float64(val), true case int64: return float64(val), true } return 0, false } // sortByScore sorts chunks by _score descending and limits func sortByScore(chunks []map[string]interface{}, limit int) []map[string]interface{} { if len(chunks) == 0 { return chunks } // Sort by _score descending sort.Slice(chunks, func(i, j int) bool { scoreI := getChunkScore(chunks[i]) scoreJ := getChunkScore(chunks[j]) return scoreI > scoreJ }) // Limit if len(chunks) > limit && limit > 0 { chunks = chunks[:limit] } return chunks } // getChunkScore extracts the score from a chunk func getChunkScore(chunk map[string]interface{}) float64 { if v, ok := chunk["_score"].(float64); ok { return v } if v, ok := chunk["SCORE"].(float64); ok { return v } if v, ok := chunk["SIMILARITY"].(float64); ok { return v } return 0.0 }