mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-20 14:41:05 +08:00
Refact functions in engine in GO (#14981)
### What problem does this PR solve? Refact functions in engine in GO ### Type of change - [x] Refactoring
This commit is contained in:
2038
internal/engine/infinity/chunk.go
Normal file
2038
internal/engine/infinity/chunk.go
Normal file
File diff suppressed because it is too large
Load Diff
@@ -25,101 +25,58 @@ import (
|
||||
"strings"
|
||||
|
||||
infinity "github.com/infiniflow/infinity-go-sdk"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// Delete deletes rows from either a dataset table or metadata table.
|
||||
// If indexName starts with "ragflow_doc_meta_", it's a metadata table.
|
||||
// Otherwise, it's a dataset table: {indexName}_{datasetID}
|
||||
func (e *infinityEngine) Delete(ctx context.Context, condition map[string]interface{}, indexName string, datasetID string) (int64, error) {
|
||||
var tableName string
|
||||
if strings.HasPrefix(indexName, "ragflow_doc_meta_") {
|
||||
tableName = indexName
|
||||
} else {
|
||||
tableName = fmt.Sprintf("%s_%s", indexName, datasetID)
|
||||
// dropTable drops a table from Infinity
|
||||
func (e *infinityEngine) dropTable(ctx context.Context, tableName string) error {
|
||||
if tableName == "" {
|
||||
return fmt.Errorf("table name cannot be empty")
|
||||
}
|
||||
|
||||
// Check if table exists
|
||||
exists, err := e.tableExists(ctx, tableName)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to check table existence: %w", err)
|
||||
}
|
||||
if !exists {
|
||||
return fmt.Errorf("table '%s' does not exist", tableName)
|
||||
}
|
||||
|
||||
db, err := e.client.conn.GetDatabase(e.client.dbName)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("failed to get database: %w", err)
|
||||
return fmt.Errorf("failed to get database: %w", err)
|
||||
}
|
||||
|
||||
table, err := db.GetTable(tableName)
|
||||
_, err = db.DropTable(tableName, infinity.ConflictTypeError)
|
||||
if err != nil {
|
||||
common.Warn(fmt.Sprintf("Table %s does not exist, skipping delete", tableName))
|
||||
return 0, nil
|
||||
return fmt.Errorf("failed to drop table: %w", err)
|
||||
}
|
||||
|
||||
// Get table columns for building filter
|
||||
clmns := make(map[string]struct {
|
||||
Type string
|
||||
Default interface{}
|
||||
})
|
||||
colsResp, err := table.ShowColumns()
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("failed to get columns: %w", err)
|
||||
}
|
||||
result, ok := colsResp.(*infinity.QueryResult)
|
||||
if ok {
|
||||
if nameArr, ok := result.Data["name"]; ok {
|
||||
if typeArr, ok := result.Data["type"]; ok {
|
||||
if defArr, ok := result.Data["default"]; ok {
|
||||
for i := 0; i < len(nameArr); i++ {
|
||||
colName, _ := nameArr[i].(string)
|
||||
colType, _ := typeArr[i].(string)
|
||||
var colDefault interface{}
|
||||
if i < len(defArr) {
|
||||
colDefault = defArr[i]
|
||||
}
|
||||
clmns[colName] = struct {
|
||||
Type string
|
||||
Default interface{}
|
||||
}{colType, colDefault}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Build filter from condition
|
||||
filter := buildFilterFromCondition(condition, clmns)
|
||||
|
||||
delResp, err := table.Delete(filter)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("failed to delete: %w", err)
|
||||
}
|
||||
|
||||
return delResp.DeletedRows, nil
|
||||
}
|
||||
|
||||
// DropTable deletes a table/index
|
||||
func (e *infinityEngine) DropTable(ctx context.Context, indexName string) error {
|
||||
db, err := e.client.conn.GetDatabase(e.client.dbName)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to get database: %w", err)
|
||||
}
|
||||
|
||||
_, err = db.DropTable(indexName, infinity.ConflictTypeIgnore)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to drop table: %w", err)
|
||||
}
|
||||
common.Info("Infinity dropped table", zap.String("tableName", tableName))
|
||||
return nil
|
||||
}
|
||||
|
||||
// TableExists checks if table/index exists
|
||||
func (e *infinityEngine) TableExists(ctx context.Context, indexName string) (bool, error) {
|
||||
db, err := e.client.conn.GetDatabase(e.client.dbName)
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("Failed to get database: %w", err)
|
||||
// tableExists checks if a table exists in Infinity
|
||||
func (e *infinityEngine) tableExists(ctx context.Context, tableName string) (bool, error) {
|
||||
if tableName == "" {
|
||||
return false, fmt.Errorf("table name cannot be empty")
|
||||
}
|
||||
|
||||
_, err = db.GetTable(indexName)
|
||||
db, err := e.client.conn.GetDatabase(e.client.dbName)
|
||||
if err != nil {
|
||||
// Check if error is "table not found"
|
||||
errLower := strings.ToLower(err.Error())
|
||||
if strings.Contains(errLower, "not found") || strings.Contains(errLower, "notexist") || strings.Contains(errLower, "doesn't exist") {
|
||||
return false, fmt.Errorf("failed to get database: %w", err)
|
||||
}
|
||||
|
||||
// Try to get the table - if it exists, no error
|
||||
_, err = db.GetTable(tableName)
|
||||
if err != nil {
|
||||
errMsg := strings.ToLower(err.Error())
|
||||
if strings.Contains(errMsg, "not found") || strings.Contains(errMsg, "doesn't exist") {
|
||||
return false, nil
|
||||
}
|
||||
return false, err
|
||||
return false, fmt.Errorf("failed to check table existence: %w", err)
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
@@ -335,3 +292,18 @@ func (e *infinityEngine) columnExists(table *infinity.Table, columnName string)
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// buildChunkTableName returns the chunk table name for a dataset
|
||||
// Skill Table: table name is just baseName (e.g., "skill_abc123_def456")
|
||||
// Regular chunk Table: table name is {baseName}_{datasetID}
|
||||
func buildChunkTableName(baseName, datasetID string) string {
|
||||
if datasetID == "skill" {
|
||||
return baseName
|
||||
}
|
||||
return fmt.Sprintf("%s_%s", baseName, datasetID)
|
||||
}
|
||||
|
||||
// buildMetadataTableName returns the metadata table name for a tenant
|
||||
func buildMetadataTableName(tenantID string) string {
|
||||
return fmt.Sprintf("ragflow_doc_meta_%s", tenantID)
|
||||
}
|
||||
|
||||
@@ -1,655 +0,0 @@
|
||||
//
|
||||
// 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 infinity
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"ragflow/internal/common"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/utility"
|
||||
|
||||
infinity "github.com/infiniflow/infinity-go-sdk"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// CreateDataset creates a table in Infinity
|
||||
// indexName is the table name prefix (e.g., "ragflow_<tenant_id>")
|
||||
// The full table name is built as "{indexName}_{datasetID}"
|
||||
// For skill index (datasetID="skill"), tableName is just indexName and uses skill_infinity_mapping.json
|
||||
func (e *infinityEngine) CreateDataset(ctx context.Context, indexName, datasetID string, vectorSize int, parserID string) error {
|
||||
vecSize := vectorSize
|
||||
|
||||
// Determine table name and mapping file based on index type
|
||||
var tableName string
|
||||
var mappingFile string
|
||||
|
||||
if datasetID == "skill" {
|
||||
// Skill index: table name is just indexName (e.g., "skill_abc123_def456")
|
||||
tableName = indexName
|
||||
mappingFile = "skill_infinity_mapping.json"
|
||||
common.Info("Creating skill index table", zap.String("tableName", tableName), zap.String("mappingFile", mappingFile))
|
||||
} else {
|
||||
// Regular document index: table name is {indexName}_{datasetID}
|
||||
tableName = fmt.Sprintf("%s_%s", indexName, datasetID)
|
||||
mappingFile = e.mappingFileName
|
||||
common.Info("Creating regular index table", zap.String("tableName", tableName), zap.String("mappingFile", mappingFile))
|
||||
}
|
||||
|
||||
// Use configured schema
|
||||
fpMapping := filepath.Join(utility.GetProjectRoot(), "conf", mappingFile)
|
||||
|
||||
schemaData, err := os.ReadFile(fpMapping)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to read mapping file: %w", err)
|
||||
}
|
||||
|
||||
var schema orderedFields
|
||||
if err := json.Unmarshal(schemaData, &schema); err != nil {
|
||||
return fmt.Errorf("Failed to parse mapping file: %w", err)
|
||||
}
|
||||
|
||||
// Get database
|
||||
db, err := e.client.conn.GetDatabase(e.client.dbName)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to get database: %w", err)
|
||||
}
|
||||
|
||||
// Determine vector column name
|
||||
vectorColName := fmt.Sprintf("q_%d_vec", vecSize)
|
||||
|
||||
// Check if table already exists
|
||||
exists, err := e.TableExists(ctx, tableName)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to check if table exists: %w", err)
|
||||
}
|
||||
|
||||
var table *infinity.Table
|
||||
if exists {
|
||||
// Table exists, open it and check if vector column needs to be added
|
||||
common.Info("Table already exists, checking for vector column", zap.String("tableName", tableName))
|
||||
table, err = db.GetTable(tableName)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to open existing table %s: %w", tableName, err)
|
||||
}
|
||||
|
||||
// Check if vector column exists (for embedding model changes)
|
||||
colExists, err := e.columnExists(table, vectorColName)
|
||||
if err != nil {
|
||||
common.Warn("Failed to check column existence", zap.String("column", vectorColName), zap.Error(err))
|
||||
}
|
||||
|
||||
// Add new vector column if it doesn't exist (handles embedding model change)
|
||||
if !colExists {
|
||||
common.Info("Adding new vector column for embedding model change", zap.String("column", vectorColName), zap.Int("size", vecSize))
|
||||
addColSchema := infinity.TableSchema{
|
||||
&infinity.ColumnDefinition{
|
||||
Name: vectorColName,
|
||||
DataType: fmt.Sprintf("vector,%d,float", vecSize),
|
||||
},
|
||||
}
|
||||
if _, err := table.AddColumns(addColSchema); err != nil {
|
||||
common.Error("Failed to add vector column "+vectorColName, err)
|
||||
return fmt.Errorf("Failed to add vector column %s: %w", vectorColName, err)
|
||||
}
|
||||
common.Info("Successfully added vector column", zap.String("column", vectorColName))
|
||||
}
|
||||
} else {
|
||||
// Table doesn't exist, create it with vector column in the initial schema
|
||||
common.Info(fmt.Sprintf("Creating table with vector column: %s with dimension %d", vectorColName, vecSize))
|
||||
|
||||
// Build column definitions (preserving JSON order)
|
||||
var columns infinity.TableSchema
|
||||
for _, fieldName := range schema.Keys {
|
||||
fieldInfo := schema.Fields[fieldName]
|
||||
col := infinity.ColumnDefinition{
|
||||
Name: fieldName,
|
||||
DataType: fieldInfo.Type,
|
||||
Default: fieldInfo.Default,
|
||||
// Comment: fieldInfo.Comment,
|
||||
}
|
||||
columns = append(columns, &col)
|
||||
}
|
||||
|
||||
// Add vector column
|
||||
columns = append(columns, &infinity.ColumnDefinition{
|
||||
Name: vectorColName,
|
||||
DataType: fmt.Sprintf("vector,%d,float", vecSize),
|
||||
})
|
||||
|
||||
// Add chunk_data column for table parser
|
||||
if parserID == "table" {
|
||||
columns = append(columns, &infinity.ColumnDefinition{
|
||||
Name: "chunk_data",
|
||||
DataType: "json",
|
||||
Default: "{}",
|
||||
})
|
||||
}
|
||||
|
||||
// Create table
|
||||
table, err = db.CreateTable(tableName, columns, infinity.ConflictTypeIgnore)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to create table: %w", err)
|
||||
}
|
||||
common.Debug("Infinity created table", zap.String("tableName", tableName))
|
||||
}
|
||||
|
||||
// Create HNSW index on vector column with unique name based on vector size
|
||||
// Use unique index name to avoid conflict when embedding model changes
|
||||
vectorIndexName := fmt.Sprintf("q_%d_vec_idx", vecSize)
|
||||
_, err = table.CreateIndex(
|
||||
vectorIndexName,
|
||||
infinity.NewIndexInfo(vectorColName, infinity.IndexTypeHnsw, map[string]string{
|
||||
"M": "16",
|
||||
"ef_construction": "50",
|
||||
"metric": "cosine",
|
||||
"encode": "lvq",
|
||||
}),
|
||||
infinity.ConflictTypeIgnore,
|
||||
"",
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to create HNSW index %s: %w", vectorIndexName, err)
|
||||
}
|
||||
common.Info("Created vector index", zap.String("indexName", vectorIndexName), zap.String("column", vectorColName))
|
||||
|
||||
// Create full-text indexes for varchar fields with analyzers
|
||||
for _, fieldName := range schema.Keys {
|
||||
fieldInfo := schema.Fields[fieldName]
|
||||
if fieldInfo.Type != "varchar" || fieldInfo.Analyzer == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
analyzers := []string{}
|
||||
switch a := fieldInfo.Analyzer.(type) {
|
||||
case string:
|
||||
analyzers = []string{a}
|
||||
case []interface{}:
|
||||
for _, v := range a {
|
||||
if s, ok := v.(string); ok {
|
||||
analyzers = append(analyzers, s)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for _, analyzer := range analyzers {
|
||||
indexNameFt := fmt.Sprintf("ft_%s_%s",
|
||||
regexp.MustCompile(`[^a-zA-Z0-9]`).ReplaceAllString(fieldName, "_"),
|
||||
regexp.MustCompile(`[^a-zA-Z0-9]`).ReplaceAllString(analyzer, "_"),
|
||||
)
|
||||
_, err = table.CreateIndex(
|
||||
indexNameFt,
|
||||
infinity.NewIndexInfo(fieldName, infinity.IndexTypeFullText, map[string]string{"ANALYZER": analyzer}),
|
||||
infinity.ConflictTypeIgnore,
|
||||
"",
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to create fulltext index %s: %w", indexNameFt, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Create secondary indexes for fields with index_type
|
||||
for _, fieldName := range schema.Keys {
|
||||
fieldInfo := schema.Fields[fieldName]
|
||||
if fieldInfo.IndexType == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
indexTypeStr := ""
|
||||
params := map[string]string{}
|
||||
|
||||
switch it := fieldInfo.IndexType.(type) {
|
||||
case string:
|
||||
indexTypeStr = it
|
||||
case map[string]interface{}:
|
||||
if t, ok := it["type"].(string); ok {
|
||||
indexTypeStr = t
|
||||
}
|
||||
if card, ok := it["cardinality"].(string); ok {
|
||||
params["cardinality"] = card
|
||||
}
|
||||
}
|
||||
|
||||
if indexTypeStr == "secondary" {
|
||||
indexNameSec := fmt.Sprintf("sec_%s", fieldName)
|
||||
_, err = table.CreateIndex(
|
||||
indexNameSec,
|
||||
infinity.NewIndexInfo(fieldName, infinity.IndexTypeSecondary, params),
|
||||
infinity.ConflictTypeIgnore,
|
||||
"",
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to create secondary index %s: %w", indexNameSec, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
_ = table // suppress unused variable warning
|
||||
return nil
|
||||
}
|
||||
|
||||
// InsertDataset inserts chunks into a dataset table
|
||||
// Table name format: {tableNamePrefix}_{knowledgebaseID}
|
||||
// Auto-create the table if it doesn't exist
|
||||
// Delete existing rows with matching IDs before insert
|
||||
func (e *infinityEngine) InsertDataset(ctx context.Context, chunks []map[string]interface{}, tableNamePrefix string, knowledgebaseID string) ([]string, error) {
|
||||
tableName := fmt.Sprintf("%s_%s", tableNamePrefix, knowledgebaseID)
|
||||
common.Info("InfinityConnection.InsertDataset called", zap.String("tableName", tableName), zap.Int("chunkCount", len(chunks)))
|
||||
|
||||
db, err := e.client.conn.GetDatabase(e.client.dbName)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Failed to get database: %w", err)
|
||||
}
|
||||
|
||||
table, err := db.GetTable(tableName)
|
||||
if err != nil {
|
||||
// Table doesn't exist, try to create it
|
||||
errMsg := strings.ToLower(err.Error())
|
||||
if !strings.Contains(errMsg, "not found") && !strings.Contains(errMsg, "doesn't exist") {
|
||||
return nil, fmt.Errorf("Failed to get table %s: %w", tableName, err)
|
||||
}
|
||||
|
||||
// Infer vector size from chunks
|
||||
vectorSize := 0
|
||||
vectorPattern := regexp.MustCompile(`q_(\d+)_vec`)
|
||||
for _, chunk := range chunks {
|
||||
for key := range chunk {
|
||||
matches := vectorPattern.FindStringSubmatch(key)
|
||||
if len(matches) >= 2 {
|
||||
vectorSize, _ = strconv.Atoi(matches[1])
|
||||
break
|
||||
}
|
||||
}
|
||||
if vectorSize > 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
if vectorSize == 0 {
|
||||
return nil, fmt.Errorf("cannot infer vector size from chunks")
|
||||
}
|
||||
|
||||
// Determine parser_id from chunk structure
|
||||
parserID := ""
|
||||
if chunkData, ok := chunks[0]["chunk_data"].(map[string]interface{}); ok && chunkData != nil {
|
||||
parserID = "table"
|
||||
}
|
||||
|
||||
// Create table
|
||||
if err := e.CreateDataset(ctx, tableNamePrefix, knowledgebaseID, vectorSize, parserID); err != nil {
|
||||
return nil, fmt.Errorf("Failed to create table: %w", err)
|
||||
}
|
||||
|
||||
table, err = db.GetTable(tableName)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Failed to get table after creation: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Get embedding columns and their sizes
|
||||
var embeddingCols [][2]interface{}
|
||||
colsResp, err := table.ShowColumns()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Failed to get columns: %w", err)
|
||||
}
|
||||
result, ok := colsResp.(*infinity.QueryResult)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unexpected response type: %T", colsResp)
|
||||
}
|
||||
|
||||
// ShowColumns returns a result set where Data contains arrays of column values
|
||||
re := regexp.MustCompile(`Embedding\([a-z]+,(\d+)\)`)
|
||||
if nameArr, ok := result.Data["name"]; ok {
|
||||
if typeArr, ok := result.Data["type"]; ok {
|
||||
for i := 0; i < len(nameArr); i++ {
|
||||
colName, _ := nameArr[i].(string)
|
||||
colType, _ := typeArr[i].(string)
|
||||
matches := re.FindStringSubmatch(colType)
|
||||
if len(matches) >= 2 {
|
||||
size, _ := strconv.Atoi(matches[1])
|
||||
embeddingCols = append(embeddingCols, [2]interface{}{colName, size})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Transform chunks using helper function
|
||||
insertChunks := make([]map[string]interface{}, len(chunks))
|
||||
for i, chunk := range chunks {
|
||||
insertChunks[i] = TransformChunkFields(chunk, embeddingCols)
|
||||
}
|
||||
|
||||
// Delete existing rows with matching IDs
|
||||
if len(insertChunks) > 0 {
|
||||
idList := make([]string, len(insertChunks))
|
||||
for i, chunk := range insertChunks {
|
||||
idList[i] = fmt.Sprintf("'%v'", chunk["id"])
|
||||
}
|
||||
filter := fmt.Sprintf("id IN (%s)", strings.Join(idList, ", "))
|
||||
common.Debug(fmt.Sprintf("Deleting existing rows with filter: %s", filter))
|
||||
delResp, delErr := table.Delete(filter)
|
||||
if delErr != nil {
|
||||
common.Warn(fmt.Sprintf("Failed to delete existing rows: %v", delErr))
|
||||
} else {
|
||||
common.Info(fmt.Sprintf("Deleted %d existing rows", delResp.DeletedRows))
|
||||
}
|
||||
}
|
||||
|
||||
// Insert chunks to dataset
|
||||
_, err = table.Insert(insertChunks)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Failed to insert chunks to dataset: %w", err)
|
||||
}
|
||||
|
||||
common.Info("InfinityConnection.InsertDataset result", zap.String("tableName", tableName), zap.Int("count", len(insertChunks)))
|
||||
return []string{}, nil
|
||||
}
|
||||
|
||||
// UpdateDataset updates chunks in a dataset table
|
||||
// Table name format: {tableNamePrefix}_{knowledgebaseID}
|
||||
func (e *infinityEngine) UpdateDataset(ctx context.Context, condition map[string]interface{}, newValue map[string]interface{}, tableNamePrefix string, knowledgebaseID string) error {
|
||||
tableName := fmt.Sprintf("%s_%s", tableNamePrefix, knowledgebaseID)
|
||||
common.Info("InfinityConnection.UpdateDataset called", zap.String("tableName", tableName), zap.Any("condition", condition))
|
||||
|
||||
db, err := e.client.conn.GetDatabase(e.client.dbName)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to get database: %w", err)
|
||||
}
|
||||
|
||||
table, err := db.GetTable(tableName)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to get table %s: %w", tableName, err)
|
||||
}
|
||||
|
||||
// Get table columns
|
||||
clmns := make(map[string]struct {
|
||||
Type string
|
||||
Default interface{}
|
||||
})
|
||||
colsResp, err := table.ShowColumns()
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to get columns: %w", err)
|
||||
}
|
||||
result, ok := colsResp.(*infinity.QueryResult)
|
||||
if ok {
|
||||
if nameArr, ok := result.Data["name"]; ok {
|
||||
if typeArr, ok := result.Data["type"]; ok {
|
||||
if defArr, ok := result.Data["default"]; ok {
|
||||
for i := 0; i < len(nameArr); i++ {
|
||||
colName, _ := nameArr[i].(string)
|
||||
colType, _ := typeArr[i].(string)
|
||||
var colDefault interface{}
|
||||
if i < len(defArr) {
|
||||
colDefault = defArr[i]
|
||||
}
|
||||
clmns[colName] = struct {
|
||||
Type string
|
||||
Default interface{}
|
||||
}{colType, colDefault}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Build filter string from condition
|
||||
filter := buildFilterFromCondition(condition, clmns)
|
||||
|
||||
// Process remove operation first
|
||||
removeValue := make(map[string]interface{})
|
||||
if removeData, ok := newValue["remove"].(map[string]interface{}); ok {
|
||||
removeValue = removeData
|
||||
}
|
||||
delete(newValue, "remove")
|
||||
|
||||
// Transform new_value fields using helper function (no embeddings needed for update)
|
||||
transformed := TransformChunkFields(newValue, nil)
|
||||
for k, v := range transformed {
|
||||
newValue[k] = v
|
||||
}
|
||||
|
||||
// Remove original fields that were transformed (they're now in transformed with new names/types)
|
||||
// Also remove intermediate token fields that shouldn't be stored in Infinity
|
||||
// This must match Python's delete list in infinity_conn.py
|
||||
for _, key := range []string{"docnm_kwd", "title_tks", "title_sm_tks", "important_kwd", "important_tks",
|
||||
"content_with_weight", "content_ltks", "content_sm_ltks", "authors_tks", "authors_sm_tks",
|
||||
"question_kwd", "question_tks"} {
|
||||
delete(newValue, key)
|
||||
}
|
||||
|
||||
// Handle remove operations if any
|
||||
if len(removeValue) > 0 {
|
||||
colToRemove := make([]string, 0, len(removeValue))
|
||||
for k := range removeValue {
|
||||
colToRemove = append(colToRemove, k)
|
||||
}
|
||||
colToRemove = append(colToRemove, "id")
|
||||
|
||||
// Query rows to be updated
|
||||
queryResult, err := table.Output(colToRemove).Filter(filter).ToResult()
|
||||
if err != nil {
|
||||
common.Warn(fmt.Sprintf("Failed to query rows for remove operation: %v", err))
|
||||
} else {
|
||||
qr, ok := queryResult.(*infinity.QueryResult)
|
||||
if ok && len(qr.Data) > 0 {
|
||||
// Get the id column and columns to remove
|
||||
idCol := qr.Data["id"]
|
||||
removeOpt := make(map[string]map[string][]string) // column -> value -> [ids]
|
||||
|
||||
for colName, colData := range qr.Data {
|
||||
if colName == "id" {
|
||||
continue
|
||||
}
|
||||
removeVal := removeValue[colName]
|
||||
for i, id := range idCol {
|
||||
if i < len(colData) {
|
||||
existingVal := colData[i]
|
||||
if removeStr, ok := removeVal.(string); ok {
|
||||
// Split existing value by ### and remove the target value
|
||||
if existingStr, ok := existingVal.(string); ok {
|
||||
parts := strings.Split(existingStr, "###")
|
||||
var newParts []string
|
||||
for _, p := range parts {
|
||||
if p != removeStr {
|
||||
newParts = append(newParts, p)
|
||||
}
|
||||
}
|
||||
if len(newParts) != len(parts) {
|
||||
idStr := fmt.Sprintf("%v", id)
|
||||
if removeOpt[colName] == nil {
|
||||
removeOpt[colName] = make(map[string][]string)
|
||||
}
|
||||
removeOpt[colName][strings.Join(newParts, "###")] = append(removeOpt[colName][strings.Join(newParts, "###")], idStr)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Execute remove updates
|
||||
for colName, valueToIDs := range removeOpt {
|
||||
for newVal, ids := range valueToIDs {
|
||||
idFilter := filter + " AND id IN (" + strings.Join(ids, ", ") + ")"
|
||||
common.Info(fmt.Sprintf("INFINITY remove update: table=%s, idFilter=%s, column=%s, newValue=%v", tableName, idFilter, colName, newVal))
|
||||
_, err := table.Update(idFilter, map[string]interface{}{colName: newVal})
|
||||
if err != nil {
|
||||
common.Warn(fmt.Sprintf("Failed to remove value from column %s: %v", colName, err))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Execute the main update
|
||||
common.Info(fmt.Sprintf("INFINITY update: table=%s, filter=%s, newValue=%v", tableName, filter, newValue))
|
||||
_, err = table.Update(filter, newValue)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to update chunks: %w", err)
|
||||
}
|
||||
|
||||
common.Info("InfinityConnection.UpdateDataset completes", zap.String("tableName", tableName))
|
||||
return nil
|
||||
}
|
||||
|
||||
// TransformChunkFields transforms chunk field name for insert/update
|
||||
// It handles field name conversions and value transformations:
|
||||
// - docnm_kwd -> docnm
|
||||
// - title_kwd/title_sm_tks -> docnm (if docnm_kwd not set)
|
||||
// - important_kwd -> important_keywords (+ important_kwd_empty_count)
|
||||
// - content_with_weight/content_ltks/content_sm_ltks -> content
|
||||
// - authors_tks/authors_sm_tks -> authors
|
||||
// - question_kwd -> questions (joined with \n), question_tks -> questions (if question_kwd not set)
|
||||
// - kb_id: list -> str (first element)
|
||||
// - position_int: list -> hex_joined string
|
||||
// - page_num_int, top_int: list -> hex string
|
||||
// - *_feas fields -> JSON string
|
||||
// - keyword fields with list values -> ### joined string
|
||||
// - chunk_data: dict -> JSON string
|
||||
// - Missing embeddings filled with zeros if embeddingCols provided
|
||||
func TransformChunkFields(chunk map[string]interface{}, embeddingCols [][2]interface{}) map[string]interface{} {
|
||||
d := make(map[string]interface{})
|
||||
|
||||
for k, v := range chunk {
|
||||
switch k {
|
||||
case "docnm_kwd":
|
||||
d["docnm"] = v
|
||||
case "title_kwd":
|
||||
if _, exists := chunk["docnm_kwd"]; !exists {
|
||||
d["docnm"] = utility.ConvertToString(v)
|
||||
}
|
||||
case "title_sm_tks":
|
||||
if _, exists := chunk["docnm_kwd"]; !exists {
|
||||
d["docnm"] = utility.ConvertToString(v)
|
||||
}
|
||||
case "important_kwd":
|
||||
if list, ok := v.([]interface{}); ok {
|
||||
emptyCount := 0
|
||||
tokens := make([]string, 0)
|
||||
for _, item := range list {
|
||||
if str, ok := item.(string); ok {
|
||||
if str == "" {
|
||||
emptyCount++
|
||||
} else {
|
||||
tokens = append(tokens, str)
|
||||
}
|
||||
}
|
||||
}
|
||||
d["important_keywords"] = strings.Join(tokens, ",")
|
||||
d["important_kwd_empty_count"] = emptyCount
|
||||
} else {
|
||||
d["important_keywords"] = utility.ConvertToString(v)
|
||||
}
|
||||
case "important_tks":
|
||||
if _, exists := chunk["important_kwd"]; !exists {
|
||||
d["important_keywords"] = v
|
||||
}
|
||||
case "content_with_weight":
|
||||
d["content"] = v
|
||||
case "content_ltks":
|
||||
if _, exists := chunk["content_with_weight"]; !exists {
|
||||
d["content"] = v
|
||||
}
|
||||
case "content_sm_ltks":
|
||||
if _, exists := chunk["content_with_weight"]; !exists {
|
||||
d["content"] = v
|
||||
}
|
||||
case "authors_tks":
|
||||
d["authors"] = v
|
||||
case "authors_sm_tks":
|
||||
if _, exists := chunk["authors_tks"]; !exists {
|
||||
d["authors"] = v
|
||||
}
|
||||
case "question_kwd":
|
||||
d["questions"] = strings.Join(utility.ConvertToStringSlice(v), "\n")
|
||||
case "tag_kwd":
|
||||
d["tag_kwd"] = strings.Join(utility.ConvertToStringSlice(v), "###")
|
||||
case "question_tks":
|
||||
if _, exists := chunk["question_kwd"]; !exists {
|
||||
d["questions"] = utility.ConvertToString(v)
|
||||
}
|
||||
case "kb_id":
|
||||
if list, ok := v.([]interface{}); ok && len(list) > 0 {
|
||||
d["kb_id"] = list[0]
|
||||
} else {
|
||||
d["kb_id"] = v
|
||||
}
|
||||
case "position_int":
|
||||
if list, ok := v.([]interface{}); ok {
|
||||
d["position_int"] = utility.ConvertPositionIntArrayToHex(list)
|
||||
} else {
|
||||
d["position_int"] = v
|
||||
}
|
||||
case "page_num_int", "top_int":
|
||||
if list, ok := v.([]interface{}); ok {
|
||||
d[k] = utility.ConvertIntArrayToHex(list)
|
||||
} else {
|
||||
d[k] = v
|
||||
}
|
||||
case "chunk_data":
|
||||
d["chunk_data"] = utility.ConvertMapToJSONString(v)
|
||||
default:
|
||||
// Check for *_feas fields
|
||||
if strings.HasSuffix(k, "_feas") {
|
||||
jsonBytes, _ := json.Marshal(v)
|
||||
d[k] = string(jsonBytes)
|
||||
} else if fieldKeyword(k) {
|
||||
// keyword fields with list values -> ### joined
|
||||
if list, ok := v.([]interface{}); ok {
|
||||
d[k] = strings.Join(utility.ConvertToStringSlice(list), "###")
|
||||
} else {
|
||||
d[k] = v
|
||||
}
|
||||
} else {
|
||||
d[k] = v
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Remove intermediate token fields
|
||||
for _, key := range []string{"docnm_kwd", "title_tks", "title_sm_tks", "important_kwd", "important_tks",
|
||||
"content_with_weight", "content_ltks", "content_sm_ltks", "authors_tks", "authors_sm_tks",
|
||||
"question_kwd", "question_tks"} {
|
||||
delete(d, key)
|
||||
}
|
||||
|
||||
// Fill missing embedding columns with zeros if embedding info provided
|
||||
for _, ec := range embeddingCols {
|
||||
name, ok1 := ec[0].(string)
|
||||
size, ok2 := ec[1].(int)
|
||||
if !ok1 || !ok2 {
|
||||
continue
|
||||
}
|
||||
if _, exists := d[name]; !exists {
|
||||
zeros := make([]float64, size)
|
||||
for i := range zeros {
|
||||
zeros[i] = 0
|
||||
}
|
||||
d[name] = zeros
|
||||
}
|
||||
}
|
||||
|
||||
return d
|
||||
}
|
||||
@@ -1,303 +0,0 @@
|
||||
//
|
||||
// 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 infinity
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"ragflow/internal/common"
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/utility"
|
||||
|
||||
infinity "github.com/infiniflow/infinity-go-sdk"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// GetChunk gets a chunk by ID
|
||||
func (e *infinityEngine) GetChunk(ctx context.Context, tableName, chunkID string, kbIDs []string) (interface{}, error) {
|
||||
if e.client == nil || e.client.conn == nil {
|
||||
return nil, fmt.Errorf("Infinity client not initialized")
|
||||
}
|
||||
|
||||
// Build list of table names to search
|
||||
var tableNames []string
|
||||
if strings.HasPrefix(tableName, "ragflow_doc_meta_") {
|
||||
tableNames = []string{tableName}
|
||||
} else {
|
||||
// Search in tables like <tableName>_<kb_id> for each kbID
|
||||
if len(kbIDs) > 0 {
|
||||
for _, kbID := range kbIDs {
|
||||
tableNames = append(tableNames, fmt.Sprintf("%s_%s", tableName, kbID))
|
||||
}
|
||||
}
|
||||
// Also try the base tableName
|
||||
tableNames = append(tableNames, tableName)
|
||||
}
|
||||
|
||||
// Try each table and collect results from all tables
|
||||
db, err := e.client.conn.GetDatabase(e.client.dbName)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get database: %w", err)
|
||||
}
|
||||
|
||||
// Collect chunks from all tables (same as Python's concat_dataframes)
|
||||
allChunks := make(map[string]map[string]interface{})
|
||||
|
||||
for _, tblName := range tableNames {
|
||||
table, err := db.GetTable(tblName)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
// Query with filter for the specific chunk ID
|
||||
filter := fmt.Sprintf("id = '%s'", chunkID)
|
||||
result, err := table.Output([]string{"*"}).Filter(filter).ToResult()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
qr, ok := result.(*infinity.QueryResult)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
if len(qr.Data) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
// Convert to chunk format
|
||||
chunks := make([]map[string]interface{}, 0)
|
||||
for colName, colData := range qr.Data {
|
||||
for i, val := range colData {
|
||||
for len(chunks) <= i {
|
||||
chunks = append(chunks, make(map[string]interface{}))
|
||||
}
|
||||
chunks[i][colName] = val
|
||||
}
|
||||
}
|
||||
|
||||
// Merge chunks into allChunks (by id), keeping first non-empty value
|
||||
for _, chunk := range chunks {
|
||||
if idVal, ok := chunk["id"].(string); ok {
|
||||
if existing, exists := allChunks[idVal]; exists {
|
||||
// Merge: keep first non-empty value for each field
|
||||
for k, v := range chunk {
|
||||
if _, has := existing[k]; !has || utility.IsEmpty(v) {
|
||||
existing[k] = v
|
||||
}
|
||||
}
|
||||
} else {
|
||||
allChunks[idVal] = chunk
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Get the chunk by chunkID
|
||||
chunk, found := allChunks[chunkID]
|
||||
if !found {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
common.Debug("infinity get chunk", zap.String("chunkID", chunkID), zap.Any("tables", tableNames))
|
||||
|
||||
// Apply field mappings (same as in GetFields)
|
||||
// docnm -> docnm_kwd, title_tks, title_sm_tks
|
||||
if val, ok := chunk["docnm"].(string); ok {
|
||||
chunk["docnm_kwd"] = val
|
||||
chunk["title_tks"] = val
|
||||
chunk["title_sm_tks"] = val
|
||||
}
|
||||
|
||||
// content -> content_with_weight, content_ltks, content_sm_ltks
|
||||
if val, ok := chunk["content"].(string); ok {
|
||||
chunk["content_with_weight"] = val
|
||||
chunk["content_ltks"] = val
|
||||
chunk["content_sm_ltks"] = val
|
||||
}
|
||||
|
||||
// important_keywords -> important_kwd (split by comma), important_tks
|
||||
if val, ok := chunk["important_keywords"].(string); ok {
|
||||
if val == "" {
|
||||
chunk["important_kwd"] = []interface{}{}
|
||||
} else {
|
||||
parts := strings.Split(val, ",")
|
||||
chunk["important_kwd"] = parts
|
||||
}
|
||||
chunk["important_tks"] = val
|
||||
} else {
|
||||
chunk["important_kwd"] = []interface{}{}
|
||||
chunk["important_tks"] = []interface{}{}
|
||||
}
|
||||
|
||||
// questions -> question_kwd (split by newline), question_tks
|
||||
if val, ok := chunk["questions"].(string); ok {
|
||||
if val == "" {
|
||||
chunk["question_kwd"] = []interface{}{}
|
||||
} else {
|
||||
parts := strings.Split(val, "\n")
|
||||
chunk["question_kwd"] = parts
|
||||
}
|
||||
chunk["question_tks"] = val
|
||||
} else {
|
||||
chunk["question_kwd"] = []interface{}{}
|
||||
chunk["question_tks"] = []interface{}{}
|
||||
}
|
||||
|
||||
if posVal, ok := chunk["position_int"].(string); ok {
|
||||
chunk["position_int"] = utility.ConvertHexToPositionIntArray(posVal)
|
||||
} else {
|
||||
chunk["position_int"] = []interface{}{}
|
||||
}
|
||||
|
||||
return chunk, nil
|
||||
}
|
||||
|
||||
// GetFields applies field mappings to chunks and returns a dict keyed by chunk ID.
|
||||
// Equivalent to Python's get_fields() in infinity_conn.py.
|
||||
// When fields is nil/empty, returns all fields from chunks.
|
||||
func GetFields(chunks []map[string]interface{}, fields []string) map[string]map[string]interface{} {
|
||||
result := make(map[string]map[string]interface{})
|
||||
if len(chunks) == 0 {
|
||||
return result
|
||||
}
|
||||
|
||||
// If fields is provided, create a set for lookup
|
||||
fieldSet := make(map[string]bool)
|
||||
for _, f := range fields {
|
||||
fieldSet[f] = true
|
||||
}
|
||||
|
||||
for _, chunk := range chunks {
|
||||
// Apply field mappings
|
||||
// docnm -> docnm_kwd, title_tks, title_sm_tks
|
||||
if val, ok := chunk["docnm"].(string); ok {
|
||||
chunk["docnm_kwd"] = val
|
||||
chunk["title_tks"] = val
|
||||
chunk["title_sm_tks"] = val
|
||||
}
|
||||
|
||||
// important_keywords -> important_kwd (split by comma), important_tks
|
||||
if val, ok := chunk["important_keywords"].(string); ok {
|
||||
if val == "" {
|
||||
chunk["important_kwd"] = []interface{}{}
|
||||
} else {
|
||||
parts := strings.Split(val, ",")
|
||||
chunk["important_kwd"] = parts
|
||||
}
|
||||
chunk["important_tks"] = val
|
||||
} else {
|
||||
chunk["important_kwd"] = []interface{}{}
|
||||
chunk["important_tks"] = []interface{}{}
|
||||
}
|
||||
|
||||
// questions -> question_kwd (split by newline), question_tks
|
||||
if val, ok := chunk["questions"].(string); ok {
|
||||
if val == "" {
|
||||
chunk["question_kwd"] = []interface{}{}
|
||||
} else {
|
||||
parts := strings.Split(val, "\n")
|
||||
chunk["question_kwd"] = parts
|
||||
}
|
||||
chunk["question_tks"] = val
|
||||
} else {
|
||||
chunk["question_kwd"] = []interface{}{}
|
||||
chunk["question_tks"] = []interface{}{}
|
||||
}
|
||||
|
||||
// content -> content_with_weight, content_ltks, content_sm_ltks
|
||||
if val, ok := chunk["content"].(string); ok {
|
||||
chunk["content_with_weight"] = val
|
||||
chunk["content_ltks"] = val
|
||||
chunk["content_sm_ltks"] = val
|
||||
}
|
||||
|
||||
// authors -> authors_tks, authors_sm_tks
|
||||
if val, ok := chunk["authors"].(string); ok {
|
||||
chunk["authors_tks"] = val
|
||||
chunk["authors_sm_tks"] = val
|
||||
}
|
||||
|
||||
// position_int: convert from hex string to array format (grouped by 5)
|
||||
if val, ok := chunk["position_int"].(string); ok {
|
||||
chunk["position_int"] = utility.ConvertHexToPositionIntArray(val)
|
||||
}
|
||||
|
||||
// Convert page_num_int and top_int from hex string to array
|
||||
for _, colName := range []string{"page_num_int", "top_int"} {
|
||||
if val, ok := chunk[colName].(string); ok && val != "" {
|
||||
chunk[colName] = utility.ConvertHexToIntArray(val)
|
||||
}
|
||||
}
|
||||
|
||||
// Post-process: convert nil/empty values to empty slices for array-like fields
|
||||
// and split _kwd fields by "###" (except knowledge_graph_kwd, docnm_kwd, important_kwd, question_kwd)
|
||||
kwdNoSplit := map[string]bool{
|
||||
"knowledge_graph_kwd": true, "docnm_kwd": true,
|
||||
"important_kwd": true, "question_kwd": true,
|
||||
}
|
||||
arrayFields := []string{
|
||||
"doc_type_kwd", "important_kwd", "important_tks", "question_tks",
|
||||
"question_kwd", "authors_tks", "authors_sm_tks", "title_tks",
|
||||
"title_sm_tks", "content_ltks", "content_sm_ltks", "tag_kwd",
|
||||
}
|
||||
for _, colName := range arrayFields {
|
||||
val, ok := chunk[colName]
|
||||
if !ok || val == nil || val == "" {
|
||||
chunk[colName] = []interface{}{}
|
||||
} else if !kwdNoSplit[colName] {
|
||||
// Split by "###" for _kwd fields
|
||||
if strVal, ok := val.(string); ok && strings.Contains(strVal, "###") {
|
||||
parts := strings.Split(strVal, "###")
|
||||
var filtered []interface{}
|
||||
for _, p := range parts {
|
||||
if p != "" {
|
||||
filtered = append(filtered, p)
|
||||
}
|
||||
}
|
||||
chunk[colName] = filtered
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Handle row_id mapping - Infinity returns "ROW_ID" but we use "row_id()"
|
||||
if val, ok := chunk["ROW_ID"]; ok {
|
||||
chunk["row_id()"] = val
|
||||
delete(chunk, "ROW_ID")
|
||||
}
|
||||
|
||||
// Build result map keyed by id
|
||||
if id, ok := chunk["id"].(string); ok {
|
||||
fieldMap := make(map[string]interface{})
|
||||
for field, value := range chunk {
|
||||
if len(fieldSet) == 0 || fieldSet[field] {
|
||||
fieldMap[field] = value
|
||||
}
|
||||
}
|
||||
result[id] = fieldMap
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// GetFields is a method wrapper for infinityEngine to satisfy DocEngine interface
|
||||
func (e *infinityEngine) GetFields(chunks []map[string]interface{}, fields []string) map[string]map[string]interface{} {
|
||||
return GetFields(chunks, fields)
|
||||
}
|
||||
@@ -22,18 +22,20 @@ import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"ragflow/internal/common"
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/utility"
|
||||
|
||||
infinity "github.com/infiniflow/infinity-go-sdk"
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/utility"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// CreateMetadata creates the document metadata table/index
|
||||
func (e *infinityEngine) CreateMetadata(ctx context.Context, indexName string) error {
|
||||
// CreateMetadataStore creates a metadata table in Infinity
|
||||
// tenantID is the tenant identifier used to build the table name
|
||||
func (e *infinityEngine) CreateMetadataStore(ctx context.Context, tenantID string) error {
|
||||
tableName := buildMetadataTableName(tenantID)
|
||||
|
||||
// Get database
|
||||
db, err := e.client.conn.GetDatabase(e.client.dbName)
|
||||
if err != nil {
|
||||
@@ -41,12 +43,12 @@ func (e *infinityEngine) CreateMetadata(ctx context.Context, indexName string) e
|
||||
}
|
||||
|
||||
// Check if table already exists
|
||||
exists, err := e.TableExists(ctx, indexName)
|
||||
exists, err := e.tableExists(ctx, tableName)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to check if table exists: %w", err)
|
||||
}
|
||||
if exists {
|
||||
return fmt.Errorf("metadata table '%s' already exists", indexName)
|
||||
return fmt.Errorf("metadata table '%s' already exists", tableName)
|
||||
}
|
||||
|
||||
// Use configured doc_meta mapping file
|
||||
@@ -69,27 +71,26 @@ func (e *infinityEngine) CreateMetadata(ctx context.Context, indexName string) e
|
||||
Name: fieldName,
|
||||
DataType: fieldInfo.Type,
|
||||
Default: fieldInfo.Default,
|
||||
// Comment: fieldInfo.Comment,
|
||||
}
|
||||
columns = append(columns, &col)
|
||||
}
|
||||
|
||||
// Create table
|
||||
_, err = db.CreateTable(indexName, columns, infinity.ConflictTypeIgnore)
|
||||
_, err = db.CreateTable(tableName, columns, infinity.ConflictTypeIgnore)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to create doc meta table: %w", err)
|
||||
}
|
||||
common.Debug("Infinity created doc meta table", zap.String("tableName", indexName))
|
||||
common.Debug("Infinity created doc meta table", zap.String("tableName", tableName))
|
||||
|
||||
// Get table for creating indexes
|
||||
table, err := db.GetTable(indexName)
|
||||
table, err := db.GetTable(tableName)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to get table: %w", err)
|
||||
}
|
||||
|
||||
// Create secondary index on id
|
||||
_, err = table.CreateIndex(
|
||||
fmt.Sprintf("idx_%s_id", indexName),
|
||||
fmt.Sprintf("idx_%s_id", tableName),
|
||||
infinity.NewIndexInfo("id", infinity.IndexTypeSecondary, nil),
|
||||
infinity.ConflictTypeIgnore,
|
||||
"",
|
||||
@@ -100,7 +101,7 @@ func (e *infinityEngine) CreateMetadata(ctx context.Context, indexName string) e
|
||||
|
||||
// Create secondary index on kb_id
|
||||
_, err = table.CreateIndex(
|
||||
fmt.Sprintf("idx_%s_kb_id", indexName),
|
||||
fmt.Sprintf("idx_%s_kb_id", tableName),
|
||||
infinity.NewIndexInfo("kb_id", infinity.IndexTypeSecondary, nil),
|
||||
infinity.ConflictTypeIgnore,
|
||||
"",
|
||||
@@ -113,11 +114,10 @@ func (e *infinityEngine) CreateMetadata(ctx context.Context, indexName string) e
|
||||
}
|
||||
|
||||
// InsertMetadata inserts document metadata into tenant's metadata table
|
||||
// Table name format: ragflow_doc_meta_{tenant_id}
|
||||
// Auto-create the table if it doesn't exist
|
||||
// Replace existing metadata with same id and kb_id
|
||||
func (e *infinityEngine) InsertMetadata(ctx context.Context, metadata []map[string]interface{}, tenantID string) ([]string, error) {
|
||||
tableName := fmt.Sprintf("ragflow_doc_meta_%s", tenantID)
|
||||
tableName := buildMetadataTableName(tenantID)
|
||||
common.Info("InfinityConnection.InsertMetadata called", zap.String("tableName", tableName), zap.Int("metaCount", len(metadata)))
|
||||
|
||||
db, err := e.client.conn.GetDatabase(e.client.dbName)
|
||||
@@ -134,7 +134,7 @@ func (e *infinityEngine) InsertMetadata(ctx context.Context, metadata []map[stri
|
||||
}
|
||||
|
||||
// Create metadata table
|
||||
if createErr := e.CreateMetadata(ctx, tableName); createErr != nil {
|
||||
if createErr := e.CreateMetadataStore(ctx, tenantID); createErr != nil {
|
||||
return nil, fmt.Errorf("Failed to create metadata table: %w", createErr)
|
||||
}
|
||||
|
||||
@@ -188,12 +188,11 @@ func (e *infinityEngine) InsertMetadata(ctx context.Context, metadata []map[stri
|
||||
}
|
||||
|
||||
// UpdateMetadata updates or inserts document metadata in tenant's metadata table.
|
||||
// If a row with the given docID and kbID exists, it merges the new metadata with existing.
|
||||
// If a row with the given docID and datasetID exists, it merges the new metadata with existing.
|
||||
// If no row exists, it inserts a new row.
|
||||
// Table name format: ragflow_doc_meta_{tenant_id}
|
||||
func (e *infinityEngine) UpdateMetadata(ctx context.Context, docID string, kbID string, metaFields map[string]interface{}, tenantID string) error {
|
||||
tableName := fmt.Sprintf("ragflow_doc_meta_%s", tenantID)
|
||||
common.Info("InfinityConnection.UpdateMetadata called", zap.String("tableName", tableName), zap.String("docID", docID), zap.String("kbID", kbID))
|
||||
func (e *infinityEngine) UpdateMetadata(ctx context.Context, docID string, datasetID string, metaFields map[string]interface{}, tenantID string) error {
|
||||
tableName := buildMetadataTableName(tenantID)
|
||||
common.Info("InfinityConnection.UpdateMetadata called", zap.String("tableName", tableName), zap.String("docID", docID), zap.String("datasetID", datasetID))
|
||||
|
||||
db, err := e.client.conn.GetDatabase(e.client.dbName)
|
||||
if err != nil {
|
||||
@@ -205,10 +204,10 @@ func (e *infinityEngine) UpdateMetadata(ctx context.Context, docID string, kbID
|
||||
return fmt.Errorf("failed to get metadata table %s: %w", tableName, err)
|
||||
}
|
||||
|
||||
// Build filter to find existing row by docID and kbID
|
||||
// Build filter to find existing row by docID and datasetID
|
||||
escapedDocID := strings.ReplaceAll(docID, "'", "''")
|
||||
escapedKbID := strings.ReplaceAll(kbID, "'", "''")
|
||||
filter := fmt.Sprintf("id = '%s' AND kb_id = '%s'", escapedDocID, escapedKbID)
|
||||
escapedDatasetID := strings.ReplaceAll(datasetID, "'", "''")
|
||||
filter := fmt.Sprintf("id = '%s' AND kb_id = '%s'", escapedDocID, escapedDatasetID)
|
||||
|
||||
// Query existing metadata using the chainable API
|
||||
queryTable := table.Output([]string{"id", "kb_id", "meta_fields"}).Filter(filter).Limit(1).Offset(0)
|
||||
@@ -271,7 +270,7 @@ func (e *infinityEngine) UpdateMetadata(ctx context.Context, docID string, kbID
|
||||
// Row doesn't exist: insert new row
|
||||
insertFields := map[string]interface{}{
|
||||
"id": docID,
|
||||
"kb_id": kbID,
|
||||
"kb_id": datasetID,
|
||||
"meta_fields": utility.ConvertMapToJSONString(metaFields),
|
||||
}
|
||||
common.Info(fmt.Sprintf("UpdateMetadata: inserting new row, table=%s, newValue=%v", tableName, insertFields))
|
||||
@@ -284,3 +283,72 @@ func (e *infinityEngine) UpdateMetadata(ctx context.Context, docID string, kbID
|
||||
common.Info("InfinityConnection.UpdateMetadata completes", zap.String("tableName", tableName), zap.String("docID", docID))
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteMetadata deletes metadata from tenant's metadata table by condition
|
||||
func (e *infinityEngine) DeleteMetadata(ctx context.Context, condition map[string]interface{}, tenantID string) (int64, error) {
|
||||
tableName := buildMetadataTableName(tenantID)
|
||||
|
||||
db, err := e.client.conn.GetDatabase(e.client.dbName)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("failed to get database: %w", err)
|
||||
}
|
||||
|
||||
table, err := db.GetTable(tableName)
|
||||
if err != nil {
|
||||
common.Warn(fmt.Sprintf("Metadata table %s does not exist, skipping delete", tableName))
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// Get table columns for building filter
|
||||
clmns := make(map[string]struct {
|
||||
Type string
|
||||
Default interface{}
|
||||
})
|
||||
colsResp, err := table.ShowColumns()
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("failed to get columns: %w", err)
|
||||
}
|
||||
result, ok := colsResp.(*infinity.QueryResult)
|
||||
if ok {
|
||||
if nameArr, ok := result.Data["name"]; ok {
|
||||
if typeArr, ok := result.Data["type"]; ok {
|
||||
if defArr, ok := result.Data["default"]; ok {
|
||||
for i := 0; i < len(nameArr); i++ {
|
||||
colName, _ := nameArr[i].(string)
|
||||
colType, _ := typeArr[i].(string)
|
||||
var colDefault interface{}
|
||||
if i < len(defArr) {
|
||||
colDefault = defArr[i]
|
||||
}
|
||||
clmns[colName] = struct {
|
||||
Type string
|
||||
Default interface{}
|
||||
}{colType, colDefault}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Build filter from condition
|
||||
filter := buildFilterFromCondition(condition, clmns)
|
||||
|
||||
delResp, err := table.Delete(filter)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("failed to delete metadata: %w", err)
|
||||
}
|
||||
|
||||
return delResp.DeletedRows, nil
|
||||
}
|
||||
|
||||
// DropMetadataStore drops a metadata table from Infinity
|
||||
func (e *infinityEngine) DropMetadataStore(ctx context.Context, tenantID string) error {
|
||||
tableName := buildMetadataTableName(tenantID)
|
||||
return e.dropTable(ctx, tableName)
|
||||
}
|
||||
|
||||
// MetadataStoreExists checks if a metadata table exists in Infinity
|
||||
func (e *infinityEngine) MetadataStoreExists(ctx context.Context, tenantID string) (bool, error) {
|
||||
tableName := buildMetadataTableName(tenantID)
|
||||
return e.tableExists(ctx, tableName)
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user