diff --git a/internal/common/task.go b/internal/common/task.go index 24132264e4..c12e21722e 100644 --- a/internal/common/task.go +++ b/internal/common/task.go @@ -17,9 +17,8 @@ package common const ( - TaskTypeIngestionTask = "ingestion_task" - TaskTypeIngestionTasklet = "ingestion_tasklet" - TaskTypeIngestionTest = "ingestion_test" + TaskTypeIngestionTask = "ingestion_task" + TaskTypeIngestionTest = "ingestion_test" ) type TaskMessage struct { diff --git a/internal/dao/database.go b/internal/dao/database.go index 6af1fd0c44..427f5960f6 100644 --- a/internal/dao/database.go +++ b/internal/dao/database.go @@ -154,8 +154,6 @@ func InitDB(migrateDB bool) error { &entity.TenantModelGroup{}, &entity.IngestionTask{}, &entity.IngestionTaskLog{}, - &entity.IngestionTasklet{}, - &entity.IngestionTaskletLog{}, &entity.FileCommit{}, &entity.FileCommitItem{}, } diff --git a/internal/dao/ingestion.go b/internal/dao/ingestion.go index 5f27b498e8..cc27860163 100644 --- a/internal/dao/ingestion.go +++ b/internal/dao/ingestion.go @@ -211,15 +211,9 @@ func (dao *IngestionTaskDAO) SetStoppingByAPIServer(taskID string) (*entity.Inge } } -type TaskletInfo struct { - TaskletID string `json:"tasklet_id"` - FilesToDelete []string `json:"files_to_delete"` -} - type TaskInfo struct { - TaskID string `json:"task_id"` - FilesToDelete []string `json:"files_to_delete"` - Tasklets []TaskletInfo `json:"tasklets"` + TaskID string `json:"task_id"` + FilesToDelete []string `json:"files_to_delete"` } func (dao *IngestionTaskDAO) RemoveByAPIServerOrAdminServer(taskID string, userID *string) (*TaskInfo, error) { @@ -264,37 +258,6 @@ func (dao *IngestionTaskDAO) RemoveByAPIServerOrAdminServer(taskID string, userI taskStatus := tasks[0].Status switch taskStatus { case common.CREATED, common.STOPPED, common.COMPLETED, common.FAILED: - // get all ingestion tasklets - var tasklets []*entity.IngestionTasklet - err = tx.Where("task_id = ?", taskID).Find(&tasklets).Error - if err != nil { - return nil, err - } - var TaskletInfos []TaskletInfo - for _, tasklet := range tasklets { - // get all ingestion tasklet log - var taskletLogs []*entity.IngestionTaskletLog - err = tx.Where("tasklet_id = ?", tasklet.ID).Find(&taskletLogs).Error - - fileMap := make(map[string]bool) - for _, taskletLog := range taskletLogs { - files, ok := taskletLog.Checkpoint["files"].([]string) - if ok { - for _, file := range files { - fileMap[file] = true - } - } - } - var filesToDelete []string - for file := range fileMap { - filesToDelete = append(filesToDelete, file) - } - TaskletInfos = append(TaskletInfos, TaskletInfo{ - TaskletID: tasklet.ID, - FilesToDelete: filesToDelete, - }) - } - // get all ingestion task log var taskLogs []*entity.IngestionTaskLog err = tx.Where("task_id = ?", taskID).Find(&taskLogs).Error @@ -324,7 +287,6 @@ func (dao *IngestionTaskDAO) RemoveByAPIServerOrAdminServer(taskID string, userI taskInfo := &TaskInfo{ TaskID: taskID, FilesToDelete: filesToDelete, - Tasklets: TaskletInfos, } committed = true return taskInfo, nil @@ -416,67 +378,3 @@ func (dao *IngestionTaskLogDAO) DeleteByTaskID(taskID string) (int64, error) { result := DB.Unscoped().Where("task_id = ?", taskID).Delete(&entity.IngestionTaskLog{}) return result.RowsAffected, result.Error } - -type IngestionTaskletDAO struct{} - -func NewIngestionTaskletDAO() *IngestionTaskletDAO { - return &IngestionTaskletDAO{} -} - -func (dao *IngestionTaskletDAO) Create(ingestionTasklet *entity.IngestionTasklet) error { - return DB.Create(ingestionTasklet).Error -} - -func (dao *IngestionTaskletDAO) UpdateStatus(taskletID, status string) error { - return DB.Model(&entity.IngestionTasklet{}).Where("id = ?", taskletID).Update("status", status).Error -} -func (dao *IngestionTaskletDAO) GetAllTasklets() ([]*entity.IngestionTasklet, error) { - var tasks []*entity.IngestionTasklet - err := DB.Find(&tasks).Error - return tasks, err -} - -func (dao *IngestionTaskletDAO) ListByUserID(userID string) ([]*entity.IngestionTasklet, error) { - var tasks []*entity.IngestionTasklet - err := DB.Where("user_id = ?", userID).Find(&tasks).Error - return tasks, err -} - -func (dao *IngestionTaskletDAO) GetByID(id string) (*entity.IngestionTasklet, error) { - var task *entity.IngestionTasklet - err := DB.Where("id = ?", id).First(&task).Error - return task, err -} - -type IngestionTaskletLogDAO struct{} - -func NewIngestionTaskletLogDAO() *IngestionTaskletLogDAO { - return &IngestionTaskletLogDAO{} -} - -func (dao *IngestionTaskletLogDAO) Create(ingestionLog *entity.IngestionTaskletLog) error { - return DB.Create(ingestionLog).Error -} - -func (dao *IngestionTaskletLogDAO) ListLogsByTaskletID(taskID string) ([]*entity.IngestionTaskletLog, error) { - var tasks []*entity.IngestionTaskletLog - err := DB.Where("task_id = ?", taskID).Find(&tasks).Error - return tasks, err -} - -func (dao *IngestionTaskletLogDAO) GetLogByLogID(logID string) (*entity.IngestionTaskletLog, error) { - var task *entity.IngestionTaskletLog - err := DB.Where("id = ?", logID).First(&task).Error - return task, err -} - -func (dao *IngestionTaskletLogDAO) LatestLogByTaskletID(taskletID string) (*entity.IngestionTaskletLog, error) { - var tasklet *entity.IngestionTaskletLog - err := DB.Where("tasklet_id = ?", taskletID).Order("create_time DESC").First(&tasklet).Error - return tasklet, err -} - -func (dao *IngestionTaskletLogDAO) DeleteByTaskletID(taskID string) (int64, error) { - result := DB.Unscoped().Where("task_id = ?", taskID).Delete(&entity.IngestionTaskletLog{}) - return result.RowsAffected, result.Error -} diff --git a/internal/entity/ingestion_task.go b/internal/entity/ingestion_task.go index 0ee5f19ca3..def89ad956 100644 --- a/internal/entity/ingestion_task.go +++ b/internal/entity/ingestion_task.go @@ -42,28 +42,3 @@ type IngestionTaskLog struct { func (IngestionTaskLog) TableName() string { return "ingestion_task_log" } - -type IngestionTasklet struct { - ID string `gorm:"column:id;primaryKey;size:32" json:"id"` - TaskID string `gorm:"column:task_id;size:32;not null;index" json:"task_id"` - Schema JSONMap `gorm:"column:schema;type:longtext" json:"schema"` - Status string `gorm:"column:status;size:32;not null;" json:"status"` - BaseModel -} - -// TableName specify table name -func (IngestionTasklet) TableName() string { - return "ingestion_tasklet" -} - -type IngestionTaskletLog struct { - ID int `gorm:"column:id;primaryKey;autoIncrement" json:"id"` - TaskletID string `gorm:"column:tasklet_id;size:32;not null;index" json:"tasklet_id"` - Checkpoint JSONMap `gorm:"column:checkpoint;type:longtext;not null" json:"checkpoint"` - BaseModel -} - -// TableName specify table name -func (IngestionTaskletLog) TableName() string { - return "ingestion_tasklet_log" -} diff --git a/internal/ingestion/pipeline/template_integration_test.go b/internal/ingestion/pipeline/template_integration_test.go index b7249876bc..5167b439f2 100644 --- a/internal/ingestion/pipeline/template_integration_test.go +++ b/internal/ingestion/pipeline/template_integration_test.go @@ -55,7 +55,7 @@ func (fixedEmbedder) Encode(texts []string) ([]componentpkg.EmbeddingResult, err func TestPipelineRun_TemplateGeneral_RealComponents(t *testing.T) { - requireTokenizerPool(t) + RequireTokenizerPool(t) templatePath := filepath.Join(repoRootFromPipelineTest(t), "agent", "templates", "ingestion_pipeline_general.json") templateBytes, err := os.ReadFile(templatePath) @@ -145,7 +145,7 @@ func TestPipelineRun_TemplateGeneral_RealComponents(t *testing.T) { } func TestPipelineRun_TemplateOne_RealComponents(t *testing.T) { - requireTokenizerPool(t) + RequireTokenizerPool(t) templatePath := filepath.Join(repoRootFromPipelineTest(t), "agent", "templates", "ingestion_pipeline_one.json") templateBytes, err := os.ReadFile(templatePath) @@ -233,7 +233,7 @@ func TestPipelineRun_TemplateOne_RealComponents(t *testing.T) { } func TestPipelineRun_TemplateOne_RealComponents_PDFDeepdocChunking(t *testing.T) { - requireTokenizerPool(t) + RequireTokenizerPool(t) t.Setenv("DEEPDOC_URL", "") t.Setenv("OSSDEEPDOC_URL", "") @@ -332,7 +332,7 @@ func TestPipelineRun_TemplateOne_RealComponents_PDFDeepdocChunking(t *testing.T) } func TestPipelineRun_TemplateManual_RealComponents(t *testing.T) { - requireTokenizerPool(t) + RequireTokenizerPool(t) templatePath := filepath.Join(repoRootFromPipelineTest(t), "agent", "templates", "ingestion_pipeline_manual.json") templateBytes, err := os.ReadFile(templatePath) @@ -429,7 +429,7 @@ func TestPipelineRun_TemplateManual_RealComponents(t *testing.T) { } func TestPipelineRun_TemplateLaws_RealComponents(t *testing.T) { - requireTokenizerPool(t) + RequireTokenizerPool(t) templatePath := filepath.Join(repoRootFromPipelineTest(t), "agent", "templates", "ingestion_pipeline_laws.json") templateBytes, err := os.ReadFile(templatePath) @@ -511,7 +511,7 @@ func TestPipelineRun_TemplateLaws_RealComponents(t *testing.T) { } func TestPipelineRun_TemplatePaper_RealComponents(t *testing.T) { - requireTokenizerPool(t) + RequireTokenizerPool(t) templatePath := filepath.Join(repoRootFromPipelineTest(t), "agent", "templates", "ingestion_pipeline_paper.json") templateBytes, err := os.ReadFile(templatePath) @@ -591,7 +591,7 @@ func TestPipelineRun_TemplatePaper_RealComponents(t *testing.T) { } func TestPipelineRun_TemplateBook_RealComponents(t *testing.T) { - requireTokenizerPool(t) + RequireTokenizerPool(t) templatePath := filepath.Join(repoRootFromPipelineTest(t), "agent", "templates", "ingestion_pipeline_book.json") templateBytes, err := os.ReadFile(templatePath) @@ -676,7 +676,7 @@ func TestPipelineRun_TemplateBook_RealComponents(t *testing.T) { } func TestPipelineRun_TemplateResume_RealComponents(t *testing.T) { - requireTokenizerPool(t) + RequireTokenizerPool(t) apiKey := os.Getenv("OPENAI_API_KEY") baseURL := os.Getenv("OPENAI_BASE_URL") model := os.Getenv("OPENAI_MODEL") @@ -772,7 +772,7 @@ func TestPipelineRun_TemplateResume_RealComponents(t *testing.T) { } func TestPipelineRun_AllIngestionTemplates_RealComponentsSmoke(t *testing.T) { - requireTokenizerPool(t) + RequireTokenizerPool(t) mem := withRealTemplateDeps(t) diff --git a/internal/ingestion/service/ingestion_service.go b/internal/ingestion/service/ingestion_service.go index e0d7023d8c..b3244e49a7 100644 --- a/internal/ingestion/service/ingestion_service.go +++ b/internal/ingestion/service/ingestion_service.go @@ -65,10 +65,8 @@ type Ingestor struct { workerWg sync.WaitGroup startOnce sync.Once - ingestionTaskDAO *dao.IngestionTaskDAO - ingestionTaskLogDAO *dao.IngestionTaskLogDAO - ingestionTaskletDAO *dao.IngestionTaskletDAO - ingestionTaskletLogDAO *dao.IngestionTaskletLogDAO + ingestionTaskDAO *dao.IngestionTaskDAO + ingestionTaskLogDAO *dao.IngestionTaskLogDAO // runDocumentTask dispatches to the migrated task handler path. // Tests may override this to verify branch routing without invoking @@ -80,20 +78,18 @@ func NewIngestor(name string, maxConcurrency int32, supportedTypes []string) *In ctx, cancel := context.WithCancel(context.Background()) id := utility.GenerateUUID() ingestor := &Ingestor{ - id: id, - name: name, - ctx: ctx, - cancel: cancel, - maxConcurrency: maxConcurrency, - supportedDocTypes: supportedTypes, - version: "1.0.0", - currentTasks: make(map[string]*taskpkg.TaskContext), - taskChan: make(chan *taskpkg.TaskContext, maxConcurrency*2), - ShutdownCh: make(chan struct{}, 1), - ingestionTaskDAO: dao.NewIngestionTaskDAO(), - ingestionTaskLogDAO: dao.NewIngestionTaskLogDAO(), - ingestionTaskletDAO: dao.NewIngestionTaskletDAO(), - ingestionTaskletLogDAO: dao.NewIngestionTaskletLogDAO(), + id: id, + name: name, + ctx: ctx, + cancel: cancel, + maxConcurrency: maxConcurrency, + supportedDocTypes: supportedTypes, + version: "1.0.0", + currentTasks: make(map[string]*taskpkg.TaskContext), + taskChan: make(chan *taskpkg.TaskContext, maxConcurrency*2), + ShutdownCh: make(chan struct{}, 1), + ingestionTaskDAO: dao.NewIngestionTaskDAO(), + ingestionTaskLogDAO: dao.NewIngestionTaskLogDAO(), } ingestor.runDocumentTask = ingestor.defaultRunDocumentTask return ingestor @@ -163,7 +159,7 @@ func (e *Ingestor) Start() error { common.Info(fmt.Sprintf("task %s is already %s", taskMessage.TaskID, task.Status)) err = taskHandle.Ack() if err != nil { - common.Error(fmt.Sprintf("error nack task %s", taskMessage.TaskID), err) + common.Error(fmt.Sprintf("error ack task %s", taskMessage.TaskID), err) return err } continue diff --git a/internal/ingestion/task/chunk_process.go b/internal/ingestion/task/chunk_process.go index 329a805512..4b746c1157 100644 --- a/internal/ingestion/task/chunk_process.go +++ b/internal/ingestion/task/chunk_process.go @@ -22,10 +22,11 @@ import ( "strings" "time" - "github.com/pkoukk/tiktoken-go" "ragflow/internal/common" "ragflow/internal/tokenizer" "ragflow/internal/utility" + + "github.com/pkoukk/tiktoken-go" ) var keywordsSplitRE = regexp.MustCompile(`[,，;；、\r\n]+`) @@ -146,7 +147,10 @@ func ProcessChunksForDataflow( ck["create_timestamp_flt"] = timestamp if _, exists := ck["id"]; !exists { - text := MustGetChunkTextString(ck, "ProcessChunksForDataflow") + text, err := MustGetChunkTextString(ck) + if err != nil { + common.Error("unexpected error", err) + } ck["id"] = ChunkID(text, docID) } diff --git a/internal/ingestion/task/chunk_process_test.go b/internal/ingestion/task/chunk_process_test.go index f486d83e33..52cc1855e6 100644 --- a/internal/ingestion/task/chunk_process_test.go +++ b/internal/ingestion/task/chunk_process_test.go @@ -259,14 +259,12 @@ func TestProcessChunksForDataflow_GeneratesID(t *testing.T) { } } -func TestProcessChunksForDataflow_PanicOnListText(t *testing.T) { +func TestProcessChunksForDataflow_NoPanicOnListText(t *testing.T) { chunks := []map[string]any{{"text": []any{"bad-shape"}}} - defer func() { - if r := recover(); r == nil { - t.Error("expected panic for list-shaped text") - } - }() - _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + res := ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + if res == nil { + t.Errorf("should return valid result") + } } func TestProcessChunksForDataflow_RemovesInternalPipelineFields(t *testing.T) { diff --git a/internal/ingestion/task/chunk_utils.go b/internal/ingestion/task/chunk_utils.go index e43e261df5..1476d9daf8 100644 --- a/internal/ingestion/task/chunk_utils.go +++ b/internal/ingestion/task/chunk_utils.go @@ -116,29 +116,34 @@ func PrepareTextsForDataflowEmbedding(chunks []map[string]any) []string { text, _ = chunk["summary"].(string) } if text == "" { - text = MustGetChunkTextString(chunk, "PrepareTextsForDataflowEmbedding") + chunkText, err := MustGetChunkTextString(chunk) + if err != nil { + common.Error("chunk[text] is not string", err) + } else { + text = chunkText + } + } + if text != "" { + texts = append(texts, text) } - texts = append(texts, text) } return texts } // MustGetChunkTextString returns chunk["text"] when it is a string. // Missing text is allowed and returns empty string. -// FIXME: remove panic before production; current panic is intentional for dev/test -// so list-shaped text payloads are surfaced immediately instead of being written -// as silent bad data. -func MustGetChunkTextString(chunk map[string]any, where string) string { +func MustGetChunkTextString(chunk map[string]any) (string, error) { val, exists := chunk["text"] if !exists || val == nil { - return "" + return "", nil } text, ok := val.(string) if ok { - return text + return text, nil } - msg := fmt.Sprintf("%s: invalid chunk text type %T, expected string, chunk=%v", where, val, chunk) - common.Error(msg, nil) - panic(msg) + msg := fmt.Sprintf("invalid chunk text type %T, expected string, chunk=%v", val, chunk) + err := fmt.Errorf("unexpected chunk text type - not string") + common.Error(msg, err) + return "", err } diff --git a/internal/ingestion/task/chunk_utils_test.go b/internal/ingestion/task/chunk_utils_test.go index 6feefc3b15..8a421e980e 100644 --- a/internal/ingestion/task/chunk_utils_test.go +++ b/internal/ingestion/task/chunk_utils_test.go @@ -260,7 +260,7 @@ func TestPrepareTexts_EmptyStringFallback(t *testing.T) { {"text": ""}, } result := PrepareTextsForDataflowEmbedding(chunks) - if result[0] != "" { + if len(result) > 0 { t.Errorf("expected empty string, got %q", result[0]) } } @@ -305,29 +305,25 @@ func TestPrepareTexts_MissingTextKey(t *testing.T) { {"other_key": "value"}, } result := PrepareTextsForDataflowEmbedding(chunks) - if result[0] != "" { + if len(result) > 0 { t.Errorf("expected empty string for missing text key, got %q", result[0]) } } -func TestPrepareTexts_PanicOnListText(t *testing.T) { +func TestPrepareTexts_NoPanicOnListText(t *testing.T) { chunks := []map[string]any{ {"text": []any{"bad-shape"}}, } - defer func() { - if r := recover(); r == nil { - t.Error("expected panic for list-shaped text") - } - }() - _ = PrepareTextsForDataflowEmbedding(chunks) + result := PrepareTextsForDataflowEmbedding(chunks) + if len(result) > 0 { + t.Errorf("expected empty string for missing text key, got %q", result[0]) + } } -func TestMustGetChunkTextString_PanicOnStringSlice(t *testing.T) { +func TestMustGetChunkTextString_NoPanicOnStringSlice(t *testing.T) { chunk := map[string]any{"text": []string{"bad-shape"}} - defer func() { - if r := recover(); r == nil { - t.Error("expected panic for []string text") - } - }() - _ = MustGetChunkTextString(chunk, "unit-test") + if _, err := MustGetChunkTextString(chunk); err == nil { + t.Errorf("expect error when chunk[text] is not string") + } + }