Implement metadata search in Infinity in GO (#13706)

### What problem does this PR solve? Add cli LIST DOCUMENTS OF DATASET quoted_string ";" LIST METADATA OF DATASETS quoted_string ("," quoted_string)* ";" LIST METADATA SUMMARY OF DATASET quoted_string (DOCUMENTS quoted_string ("," quoted_string)*)? ";" ### Type of change - [x] Refactoring
2026-09-08 10:14:35 +08:00 · 2026-03-21 18:10:00 +08:00
parent db57155b30
commit 7b86f577be
10 changed files with 1084 additions and 50 deletions
@@ -17,22 +17,36 @@
 package service

 import (
+	"encoding/json"
 	"fmt"
+	"regexp"
+	"sort"
 	"time"

 	"ragflow/internal/dao"
+	"ragflow/internal/engine"
 	"ragflow/internal/model"
+	"ragflow/internal/server"
 )

 // DocumentService document service
 type DocumentService struct {
 	documentDAO *dao.DocumentDAO
+	kbDAO       *dao.KnowledgebaseDAO
+	docEngine   engine.DocEngine
+	engineType  server.EngineType
+	metadataSvc *MetadataService
 }

 // NewDocumentService create document service
 func NewDocumentService() *DocumentService {
+	cfg := server.GetConfig()
 	return &DocumentService{
 		documentDAO: dao.NewDocumentDAO(),
+		kbDAO:       dao.NewKnowledgebaseDAO(),
+		docEngine:   engine.Get(),
+		engineType:  cfg.DocEngine.Type,
+		metadataSvc: NewMetadataService(),
 	}
 }

@@ -160,6 +174,22 @@ func (s *DocumentService) ListDocuments(page, pageSize int) ([]*DocumentResponse
 	return responses, total, nil
 }

+// ListDocumentsByKBID list documents by knowledge base ID
+func (s *DocumentService) ListDocumentsByKBID(kbID string, page, pageSize int) ([]*DocumentResponse, int64, error) {
+	offset := (page - 1) * pageSize
+	documents, total, err := s.documentDAO.ListByKBID(kbID, offset, pageSize)
+	if err != nil {
+		return nil, 0, err
+	}
+
+	responses := make([]*DocumentResponse, len(documents))
+	for i, doc := range documents {
+		responses[i] = s.toResponse(doc)
+	}
+
+	return responses, total, nil
+}
+
 // GetDocumentsByAuthorID get documents by author ID
 func (s *DocumentService) GetDocumentsByAuthorID(authorID, page, pageSize int) ([]*DocumentResponse, int64, error) {
 	offset := (page - 1) * pageSize
@@ -180,7 +210,15 @@ func (s *DocumentService) GetDocumentsByAuthorID(authorID, page, pageSize int) (
 func (s *DocumentService) toResponse(doc *model.Document) *DocumentResponse {
 	createdAt := ""
 	if doc.CreateTime != nil {
-		createdAt = time.Unix(*doc.CreateTime, 0).Format("2006-01-02 15:04:05")
+		// Check if timestamp is in milliseconds (13 digits) or seconds (10 digits)
+		var ts int64
+		if *doc.CreateTime > 1000000000000 {
+			// Milliseconds - convert to seconds
+			ts = *doc.CreateTime / 1000
+		} else {
+			ts = *doc.CreateTime
+		}
+		createdAt = time.Unix(ts, 0).Format("2006-01-02 15:04:05")
 	}
 	updatedAt := ""
 	if doc.UpdateTime != nil {
@@ -209,3 +247,413 @@ func (s *DocumentService) toResponse(doc *model.Document) *DocumentResponse {
 		UpdatedAt:       updatedAt,
 	}
 }
+
+// GetMetadataSummaryRequest request for metadata summary
+type GetMetadataSummaryRequest struct {
+	KBID   string   `json:"kb_id" binding:"required"`
+	DocIDs []string `json:"doc_ids"`
+}
+
+// GetMetadataSummaryResponse response for metadata summary
+type GetMetadataSummaryResponse struct {
+	Summary map[string]interface{} `json:"summary"`
+}
+
+// GetMetadataSummary get metadata summary for documents
+func (s *DocumentService) GetMetadataSummary(kbID string, docIDs []string) (map[string]interface{}, error) {
+	tenantID, err := s.metadataSvc.GetTenantIDByKBID(kbID)
+	if err != nil {
+		return nil, err
+	}
+
+	searchResult, err := s.metadataSvc.SearchMetadata(kbID, tenantID, docIDs, 1000)
+	if err != nil {
+		return nil, err
+	}
+
+	// Aggregate metadata from results
+	return aggregateMetadata(searchResult.Chunks), nil
+}
+
+// GetDocumentMetadataByID get metadata for a specific document
+func (s *DocumentService) GetDocumentMetadataByID(docID string) (map[string]interface{}, error) {
+	// Get document to find kb_id
+	doc, err := s.documentDAO.GetByID(docID)
+	if err != nil {
+		return nil, fmt.Errorf("document not found: %w", err)
+	}
+
+	tenantID, err := s.metadataSvc.GetTenantIDByKBID(doc.KbID)
+	if err != nil {
+		return nil, err
+	}
+
+	searchResult, err := s.metadataSvc.SearchMetadata(doc.KbID, tenantID, []string{docID}, 1)
+	if err != nil {
+		return nil, err
+	}
+
+	// Return metadata if found
+	if len(searchResult.Chunks) > 0 {
+		chunk := searchResult.Chunks[0]
+		return ExtractMetaFields(chunk)
+	}
+
+	return make(map[string]interface{}), nil
+}
+
+// GetMetadataByKBs get metadata for knowledge bases
+func (s *DocumentService) GetMetadataByKBs(kbIDs []string) (map[string]interface{}, error) {
+	if len(kbIDs) == 0 {
+		return make(map[string]interface{}), nil
+	}
+
+	searchResult, err := s.metadataSvc.SearchMetadataByKBs(kbIDs, 10000)
+	if err != nil {
+		return nil, err
+	}
+
+	flattenedMeta := make(map[string]map[string][]string)
+	numChunks := len(searchResult.Chunks)
+
+	var allMetaFields []map[string]interface{}
+	if numChunks > 1 && len(searchResult.Chunks) > 0 {
+		firstChunk := searchResult.Chunks[0]
+		if metaFieldsVal := firstChunk["meta_fields"]; metaFieldsVal != nil {
+			if v, ok := metaFieldsVal.([]byte); ok {
+				allMetaFields = ParseAllLengthPrefixedJSON(v)
+			}
+		}
+	}
+
+	for idx, chunk := range searchResult.Chunks {
+		docID, ok := ExtractDocumentID(chunk)
+		if !ok {
+			continue
+		}
+
+		var metaFields map[string]interface{}
+		var metaFieldsVal interface{}
+
+		if len(allMetaFields) > 0 && idx < len(allMetaFields) {
+			// Use pre-parsed meta_fields from concatenated data
+			metaFields = allMetaFields[idx]
+		} else {
+			// Normal case - get from chunk
+			metaFieldsVal = chunk["meta_fields"]
+			if metaFieldsVal != nil {
+				switch v := metaFieldsVal.(type) {
+				case string:
+					if err := json.Unmarshal([]byte(v), &metaFields); err != nil {
+						continue
+					}
+				case []byte:
+					// Try direct JSON parse first
+					if err := json.Unmarshal(v, &metaFields); err != nil {
+						// Try to parse as concatenated JSON objects
+						metaFields = ParseLengthPrefixedJSON(v)
+					}
+				case map[string]interface{}:
+					metaFields = v
+				default:
+					continue
+				}
+			}
+		}
+
+		if metaFields == nil {
+			continue
+		}
+
+		// Process each metadata field
+		for fieldName, fieldValue := range metaFields {
+			if fieldName == "kb_id" || fieldName == "id" {
+				continue
+			}
+
+			if _, ok := flattenedMeta[fieldName]; !ok {
+				flattenedMeta[fieldName] = make(map[string][]string)
+			}
+
+			// Handle list and single values
+			var values []interface{}
+			switch v := fieldValue.(type) {
+			case []interface{}:
+				values = v
+			default:
+				values = []interface{}{v}
+			}
+
+			for _, val := range values {
+				if val == nil {
+					continue
+				}
+				strVal := fmt.Sprintf("%v", val)
+				flattenedMeta[fieldName][strVal] = append(flattenedMeta[fieldName][strVal], docID)
+			}
+		}
+	}
+
+	// Convert to map[string]interface{} for return
+	var metaResult map[string]interface{} = make(map[string]interface{})
+	for k, v := range flattenedMeta {
+		metaResult[k] = v
+	}
+
+	return metaResult, nil
+}
+
+// valueInfo holds count and order of first appearance
+type valueInfo struct {
+	count     int
+	firstOrder int
+}
+
+// aggregateMetadata aggregates metadata from search results
+func aggregateMetadata(chunks []map[string]interface{}) map[string]interface{} {
+	// summary: map[fieldName]map[value]valueInfo
+	summary := make(map[string]map[string]valueInfo)
+	typeCounter := make(map[string]map[string]int)
+	orderCounter := 0
+
+	for _, chunk := range chunks {
+		// For metadata table, the actual metadata is in the "meta_fields" JSON field
+		// Extract it first
+		metaFieldsVal := chunk["meta_fields"]
+		if metaFieldsVal == nil {
+			continue
+		}
+
+		// Parse meta_fields - could be a string (JSON) or a map
+		var metaFields map[string]interface{}
+		switch v := metaFieldsVal.(type) {
+		case string:
+			// Parse JSON string
+			if err := json.Unmarshal([]byte(v), &metaFields); err != nil {
+				continue
+			}
+		case []byte:
+			// Handle byte slice - Infinity returns concatenated JSON objects with length prefixes
+			rawBytes := v
+
+			// Try to detect and handle length-prefixed format
+			// Format: [4-byte length][JSON][4-byte length][JSON]...
+			parsedMetaFields := make(map[string]interface{})
+			offset := 0
+			for offset < len(rawBytes) {
+				// Need at least 4 bytes for length prefix
+				if offset+4 > len(rawBytes) {
+					break
+				}
+
+				// Read 4-byte length (little-endian, not big-endian!)
+				length := uint32(rawBytes[offset]) | uint32(rawBytes[offset+1])<<8 |
+					uint32(rawBytes[offset+2])<<16 | uint32(rawBytes[offset+3])<<24
+
+				// Check if length looks valid (not too large)
+				if length > 10000 || length == 0 {
+					// Try to find next '{' from current position
+					nextBrace := -1
+					for i := offset; i < len(rawBytes) && i < offset+100; i++ {
+						if rawBytes[i] == '{' {
+							nextBrace = i
+							break
+						}
+					}
+					if nextBrace > offset {
+						// Skip to the next '{'
+						offset = nextBrace
+						continue
+					}
+					break
+				}
+
+				// Extract JSON data
+				jsonStart := offset + 4
+				jsonEnd := jsonStart + int(length)
+				if jsonEnd > len(rawBytes) {
+					jsonEnd = len(rawBytes)
+				}
+
+				jsonBytes := rawBytes[jsonStart:jsonEnd]
+
+				// Try to parse this JSON
+				var singleMeta map[string]interface{}
+				if err := json.Unmarshal(jsonBytes, &singleMeta); err == nil {
+					// Merge metadata from this document
+					for k, vv := range singleMeta {
+						if existing, ok := parsedMetaFields[k]; ok {
+							// Combine values
+							if existList, ok := existing.([]interface{}); ok {
+								if newList, ok := vv.([]interface{}); ok {
+									parsedMetaFields[k] = append(existList, newList...)
+								} else {
+									parsedMetaFields[k] = append(existList, vv)
+								}
+							} else {
+								parsedMetaFields[k] = []interface{}{existing, vv}
+							}
+						} else {
+							parsedMetaFields[k] = vv
+						}
+					}
+				}
+
+				offset = jsonEnd
+			}
+
+			// If we successfully parsed multiple JSON objects, use the merged result
+			if len(parsedMetaFields) > 0 {
+				metaFields = parsedMetaFields
+			} else {
+				// Fallback: try the original parsing method
+				startIdx := -1
+				for i, b := range rawBytes {
+					if b == '{' {
+						startIdx = i
+						break
+					}
+				}
+				if startIdx > 0 {
+					strVal := string(rawBytes[startIdx:])
+					if err := json.Unmarshal([]byte(strVal), &metaFields); err != nil {
+						metaFields = map[string]interface{}{"raw": strVal}
+					}
+				} else if err := json.Unmarshal(rawBytes, &metaFields); err != nil {
+					metaFields = map[string]interface{}{"raw": string(rawBytes)}
+				}
+			}
+		case map[string]interface{}:
+			metaFields = v
+		default:
+			continue
+		}
+
+		// Now iterate over the extracted metadata fields
+		for k, v := range metaFields {
+			// Skip nil values
+			if v == nil {
+				continue
+			}
+
+			// Determine value type
+			valueType := getMetaValueType(v)
+
+			// Track type counts
+			if valueType != "" {
+				if _, ok := typeCounter[k]; !ok {
+					typeCounter[k] = make(map[string]int)
+				}
+				typeCounter[k][valueType] = typeCounter[k][valueType] + 1
+			}
+
+			// Aggregate value counts
+			values := v
+			if v, ok := v.([]interface{}); ok {
+				values = v
+			} else {
+				values = []interface{}{v}
+			}
+
+			for _, vv := range values.([]interface{}) {
+				if vv == nil {
+					continue
+				}
+				sv := fmt.Sprintf("%v", vv)
+
+				if _, ok := summary[k]; !ok {
+					summary[k] = make(map[string]valueInfo)
+				}
+
+				if existing, ok := summary[k][sv]; ok {
+					// Already exists, just increment count
+					existing.count++
+					summary[k][sv] = existing
+				} else {
+					// First time seeing this value - record order
+					summary[k][sv] = valueInfo{count: 1, firstOrder: orderCounter}
+					orderCounter++
+				}
+			}
+		}
+	}
+
+	// Build result with type information and sorted values
+	result := make(map[string]interface{})
+	for k, v := range summary {
+		// Sort by count descending, then by firstOrder ascending (to match Python stable sort)
+		// values: [value, count, firstOrder]
+		values := make([][3]interface{}, 0, len(v))
+		for val, info := range v {
+			values = append(values, [3]interface{}{val, info.count, info.firstOrder})
+		}
+		// Use stable sort - sort by count descending, then by firstOrder
+		sort.SliceStable(values, func(i, j int) bool {
+			cntI := values[i][1].(int)
+			cntJ := values[j][1].(int)
+			if cntI != cntJ {
+				return cntI > cntJ // count descending
+			}
+			// If counts equal, use firstOrder ascending (earlier appearance first)
+			return values[i][2].(int) < values[j][2].(int)
+		})
+
+		// Determine dominant type
+		valueType := "string"
+		if typeCounts, ok := typeCounter[k]; ok {
+			maxCount := 0
+			for t, c := range typeCounts {
+				if c > maxCount {
+					maxCount = c
+					valueType = t
+				}
+			}
+		}
+
+		// Convert from [value, count, firstOrder] to [value, count] for output
+		outputValues := make([][2]interface{}, len(values))
+		for i, val := range values {
+			outputValues[i] = [2]interface{}{val[0], val[1]}
+		}
+
+		result[k] = map[string]interface{}{
+			"type":  valueType,
+			"values": outputValues,
+		}
+	}
+
+	return result
+}
+
+// getMetaValueType determines the type of a metadata value
+func getMetaValueType(value interface{}) string {
+	if value == nil {
+		return ""
+	}
+
+	switch v := value.(type) {
+	case []interface{}:
+		if len(v) > 0 {
+			return "list"
+		}
+		return ""
+	case bool:
+		return "string"
+	case int, int8, int16, int32, int64:
+		return "number"
+	case float32, float64:
+		return "number"
+	case string:
+		if isTimeString(v) {
+			return "time"
+		}
+		return "string"
+	}
+	return "string"
+}
+
+// isTimeString checks if a string is an ISO 8601 datetime
+func isTimeString(s string) bool {
+	matched, _ := regexp.MatchString(`^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}$`, s)
+	return matched
+}
@@ -0,0 +1,269 @@
+//
+//  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 service
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+
+	"ragflow/internal/dao"
+	"ragflow/internal/engine"
+	"ragflow/internal/engine/types"
+)
+
+// MetadataService provides common metadata operations
+type MetadataService struct {
+	kbDAO     *dao.KnowledgebaseDAO
+	docEngine engine.DocEngine
+}
+
+// NewMetadataService creates a new metadata service
+func NewMetadataService() *MetadataService {
+	return &MetadataService{
+		kbDAO:     dao.NewKnowledgebaseDAO(),
+		docEngine: engine.Get(),
+	}
+}
+
+// BuildMetadataIndexName constructs the metadata index name for a tenant
+func BuildMetadataIndexName(tenantID string) string {
+	return fmt.Sprintf("ragflow_doc_meta_%s", tenantID)
+}
+
+// GetTenantIDByKBID retrieves tenant ID from knowledge base ID
+func (s *MetadataService) GetTenantIDByKBID(kbID string) (string, error) {
+	kb, err := s.kbDAO.GetByID(kbID)
+	if err != nil {
+		return "", fmt.Errorf("knowledgebase not found: %w", err)
+	}
+	return kb.TenantID, nil
+}
+
+// GetTenantIDByKBIDs retrieves tenant ID from the first knowledge base ID in the list
+func (s *MetadataService) GetTenantIDByKBIDs(kbIDs []string) (string, error) {
+	if len(kbIDs) == 0 {
+		return "", fmt.Errorf("no kb_ids provided")
+	}
+	kb, err := s.kbDAO.GetByID(kbIDs[0])
+	if err != nil {
+		return "", fmt.Errorf("knowledgebase not found: %w", err)
+	}
+	return kb.TenantID, nil
+}
+
+// SearchMetadataResult holds the result of a metadata search
+type SearchMetadataResult struct {
+	IndexName string
+	Chunks    []map[string]interface{}
+}
+
+// SearchMetadata searches the metadata index with the given parameters
+func (s *MetadataService) SearchMetadata(kbID, tenantID string, docIDs []string, size int) (*SearchMetadataResult, error) {
+	indexName := BuildMetadataIndexName(tenantID)
+
+	searchReq := &types.SearchRequest{
+		IndexNames:   []string{indexName},
+		KbIDs:        []string{kbID},
+		DocIDs:       docIDs,
+		Page:         1,
+		Size:         size,
+		KeywordOnly:  true,
+	}
+
+	result, err := s.docEngine.Search(context.Background(), searchReq)
+	if err != nil {
+		return nil, fmt.Errorf("search failed: %w", err)
+	}
+
+	searchResp, ok := result.(*types.SearchResponse)
+	if !ok {
+		return nil, fmt.Errorf("invalid search response type")
+	}
+
+	return &SearchMetadataResult{
+		IndexName: indexName,
+		Chunks:    searchResp.Chunks,
+	}, nil
+}
+
+// SearchMetadataByKBs searches the metadata index for multiple knowledge bases
+func (s *MetadataService) SearchMetadataByKBs(kbIDs []string, size int) (*SearchMetadataResult, error) {
+	if len(kbIDs) == 0 {
+		return &SearchMetadataResult{Chunks: []map[string]interface{}{}}, nil
+	}
+
+	tenantID, err := s.GetTenantIDByKBIDs(kbIDs)
+	if err != nil {
+		return nil, err
+	}
+
+	indexName := BuildMetadataIndexName(tenantID)
+
+	searchReq := &types.SearchRequest{
+		IndexNames:   []string{indexName},
+		KbIDs:        kbIDs,
+		Page:         1,
+		Size:         size,
+		KeywordOnly:  true,
+	}
+
+	result, err := s.docEngine.Search(context.Background(), searchReq)
+	if err != nil {
+		return nil, fmt.Errorf("search failed: %w", err)
+	}
+
+	searchResp, ok := result.(*types.SearchResponse)
+	if !ok {
+		return nil, fmt.Errorf("invalid search response type")
+	}
+
+	return &SearchMetadataResult{
+		IndexName: indexName,
+		Chunks:    searchResp.Chunks,
+	}, nil
+}
+
+// ExtractDocumentID extracts the document ID from a chunk
+func ExtractDocumentID(chunk map[string]interface{}) (string, bool) {
+	docID, ok := chunk["id"].(string)
+	return docID, ok
+}
+
+// ExtractMetaFields extracts meta_fields from a chunk, handling different types
+func ExtractMetaFields(chunk map[string]interface{}) (map[string]interface{}, error) {
+	metaFieldsVal := chunk["meta_fields"]
+	if metaFieldsVal == nil {
+		return make(map[string]interface{}), nil
+	}
+
+	var metaFields map[string]interface{}
+	switch v := metaFieldsVal.(type) {
+	case map[string]interface{}:
+		metaFields = v
+	case string:
+		if err := json.Unmarshal([]byte(v), &metaFields); err != nil {
+			return make(map[string]interface{}), nil
+		}
+	case []byte:
+		metaFields = ParseLengthPrefixedJSON(v)
+		if metaFields == nil {
+			if err := json.Unmarshal(v, &metaFields); err != nil {
+				return make(map[string]interface{}), nil
+			}
+		}
+	default:
+		return make(map[string]interface{}), nil
+	}
+
+	return metaFields, nil
+}
+
+// ParseLengthPrefixedJSON parses Infinity's length-prefixed JSON format
+// Format: [4-byte length (little-endian)][JSON][4-byte length][JSON]...
+// Returns the FIRST valid JSON object found
+func ParseLengthPrefixedJSON(data []byte) map[string]interface{} {
+	if len(data) < 4 {
+		return nil
+	}
+
+	// Try to find the first valid JSON object by skipping length prefixes
+	offset := 0
+	for offset < len(data) {
+		// Skip non-'{' bytes
+		for offset < len(data) && data[offset] != '{' {
+			offset++
+		}
+		if offset >= len(data) {
+			break
+		}
+
+		// Try to parse JSON from current position
+		var result map[string]interface{}
+		err := json.Unmarshal(data[offset:], &result)
+		if err == nil {
+			return result
+		}
+
+		// Move forward to try next position
+		offset++
+	}
+	return nil
+}
+
+// ParseAllLengthPrefixedJSON parses Infinity's length-prefixed JSON format
+// and returns ALL JSON objects found (for cases where multiple rows are concatenated)
+// Format: [4-byte length (little-endian)][JSON][4-byte length][JSON]...
+func ParseAllLengthPrefixedJSON(data []byte) []map[string]interface{} {
+	if len(data) < 4 {
+		return nil
+	}
+
+	var results []map[string]interface{}
+	offset := 0
+
+	// Use length prefix to extract each JSON
+	for offset+4 <= len(data) {
+		// Read 4-byte length (little-endian)
+		length := uint32(data[offset]) | uint32(data[offset+1])<<8 |
+			uint32(data[offset+2])<<16 | uint32(data[offset+3])<<24
+
+		// Check if length looks reasonable
+		if length == 0 || offset+4+int(length) > len(data) {
+			// Length invalid, try to find next '{'
+			nextBrace := -1
+			for i := offset + 4; i < len(data) && i < offset+104; i++ {
+				if data[i] == '{' {
+					nextBrace = i
+					break
+				}
+			}
+			if nextBrace > offset {
+				offset = nextBrace
+				continue
+			}
+			break
+		}
+
+		// Extract JSON bytes (skip the 4-byte length prefix)
+		jsonStart := offset + 4
+		jsonEnd := jsonStart + int(length)
+		jsonBytes := data[jsonStart:jsonEnd]
+
+		var result map[string]interface{}
+		if err := json.Unmarshal(jsonBytes, &result); err == nil {
+			results = append(results, result)
+			offset = jsonEnd
+			continue
+		} else {
+			// Try to find next '{'
+			nextBrace := -1
+			for i := offset + 4; i < len(data) && i < offset+104; i++ {
+				if data[i] == '{' {
+					nextBrace = i
+					break
+				}
+			}
+			if nextBrace > offset {
+				offset = nextBrace
+				continue
+			}
+			break
+		}
+	}
+	return results
+}