diff --git a/api/apps/restful_apis/dataset_api.py b/api/apps/restful_apis/dataset_api.py index 42b2a55d8f..de0615f1da 100644 --- a/api/apps/restful_apis/dataset_api.py +++ b/api/apps/restful_apis/dataset_api.py @@ -910,21 +910,6 @@ def delete_index(tenant_id, dataset_id, index_type=None): return get_error_data_result(message="Internal server error") -@manager.route("/datasets//embedding", methods=["POST"]) # noqa: F821 -@login_required -@add_tenant_id_to_kwargs -async def run_embedding(tenant_id, dataset_id): - try: - success, result = dataset_api_service.run_embedding(dataset_id, tenant_id) - if success: - return get_result(data=result) - else: - return get_error_data_result(message=result) - except Exception as e: - logging.exception(e) - return get_error_data_result(message="Internal server error") - - @manager.route("/datasets//embedding/check", methods=["POST"]) # noqa: F821 @login_required @add_tenant_id_to_kwargs diff --git a/api/apps/services/dataset_api_service.py b/api/apps/services/dataset_api_service.py index a1a67f2ed5..9c886e5498 100644 --- a/api/apps/services/dataset_api_service.py +++ b/api/apps/services/dataset_api_service.py @@ -873,46 +873,6 @@ def delete_index(dataset_id: str, tenant_id: str, index_type: str, wipe: bool = return True, {} -def run_embedding(dataset_id: str, tenant_id: str): - """ - Run embedding for all documents in a dataset. - - :param dataset_id: dataset ID - :param tenant_id: tenant ID - :return: (success, result) or (success, error_message) - """ - if not dataset_id: - return False, 'Lack of "Dataset ID"' - - if not KnowledgebaseService.accessible(dataset_id, tenant_id): - return False, "No authorization." - - ok, kb = KnowledgebaseService.get_by_id(dataset_id) - if not ok: - return False, "Invalid Dataset ID" - - documents, _ = DocumentService.get_by_kb_id( - kb_id=dataset_id, - page_number=0, - items_per_page=0, - orderby="create_time", - desc=False, - keywords="", - run_status=[], - types=[], - suffix=[], - ) - if not documents: - return False, f"No documents in Dataset {dataset_id}" - - kb_table_num_map = {} - for doc in documents: - doc["tenant_id"] = tenant_id - DocumentService.run(tenant_id, doc, kb_table_num_map) - - return True, {"scheduled_count": len(documents)} - - def rename_tag(dataset_id: str, tenant_id: str, from_tag: str, to_tag: str): """ Rename a tag in a dataset. diff --git a/internal/handler/dataset.go b/internal/handler/dataset.go index a0504a5913..97e0317929 100644 --- a/internal/handler/dataset.go +++ b/internal/handler/dataset.go @@ -663,35 +663,6 @@ func (h *DatasetsHandler) RemoveTags(c *gin.Context) { common.SuccessWithData(c, true, "success") } -// RunEmbedding Run embedding for all documents in a dataset. -func (h *DatasetsHandler) RunEmbedding(c *gin.Context) { - user, errorCode, errorMessage := GetUser(c) - if errorCode != common.CodeSuccess { - common.ErrorWithCode(c, errorCode, errorMessage) - return - } - - userID := strings.TrimSpace(user.ID) - if userID == "" { - common.ResponseWithCodeData(c, common.CodeAuthenticationError, nil, "user_id is required") - return - } - - datasetID := strings.TrimSpace(c.Param("dataset_id")) - if datasetID == "" { - common.ResponseWithCodeData(c, common.CodeDataError, nil, "dataset_id is required") - return - } - - result, errorCode, err := h.datasetsService.RunEmbedding(userID, datasetID) - if err != nil { - common.ResponseWithCodeData(c, errorCode, nil, err.Error()) - return - } - - common.SuccessWithData(c, result, "success") -} - // CheckEmbedding Check embedding model compatibility by sampling random chunks, // re-embedding them with the new model, and computing cosine similarity. func (h *DatasetsHandler) CheckEmbedding(c *gin.Context) { diff --git a/internal/router/router.go b/internal/router/router.go index 33a1785bab..9b6dd388f5 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -365,7 +365,6 @@ func (r *Router) Setup(engine *gin.Engine) { datasets.GET("/:dataset_id/tags", r.datasetsHandler.ListTags) datasets.PUT("/:dataset_id/tags", r.datasetsHandler.RenameTag) datasets.DELETE("/:dataset_id/tags", r.datasetsHandler.RemoveTags) - datasets.POST("/:dataset_id/embedding", r.datasetsHandler.RunEmbedding) datasets.POST("/:dataset_id/embedding/check", r.datasetsHandler.CheckEmbedding) datasets.POST("/:dataset_id/documents/batch-update-status", r.documentHandler.BatchUpdateDocumentStatus) datasets.GET("/:dataset_id/index", r.datasetsHandler.TraceIndex) diff --git a/internal/service/dataset.go b/internal/service/dataset.go index ed2e288770..8e36436d48 100644 --- a/internal/service/dataset.go +++ b/internal/service/dataset.go @@ -70,7 +70,6 @@ var ( const ( // Keep the legacy worker marker in queue payloads; persisted tasks use a real document ID. graphRaptorQueueDocID = "graph_raptor_x" - maximumPageNumber = int64(100000) maximumTaskPageNumber = int64(100000000) serverQueueNamePrefix = "te" defaultEmbeddingCheckNum = 5 @@ -635,143 +634,6 @@ type embeddingCheckSample struct { QuestionKeywords []string } -// RunEmbedding runs embedding for all documents in a dataset. -func (d *DatasetService) RunEmbedding(userID, datasetID string) (map[string]interface{}, common.ErrorCode, error) { - if datasetID == "" { - return nil, common.CodeDataError, errors.New(`Lack of "Dataset ID"`) - } - if !d.kbDAO.Accessible(datasetID, userID) { - return nil, common.CodeDataError, errors.New("No authorization.") - } - - kb, err := d.kbDAO.GetByID(datasetID) - if err != nil { - if dao.IsNotFoundErr(err) { - return nil, common.CodeDataError, errors.New("Invalid Dataset ID") - } - return nil, common.CodeServerError, errors.New("Internal server error") - } - - documents, _, err := d.documentDAO.GetByKBID(datasetID) - if err != nil { - return nil, common.CodeServerError, errors.New("Internal server error") - } - if len(documents) == 0 { - return nil, common.CodeDataError, fmt.Errorf("No documents in Dataset %s", datasetID) - } - - tableDoneCountByKB := make(map[string]int64) - scheduledCount := 0 - for _, doc := range documents { - if doc == nil { - continue - } - if err := d.runEmbeddingDocument(kb, doc, tableDoneCountByKB); err != nil { - common.Warn("Failed to schedule dataset embedding document", - zap.String("datasetID", datasetID), - zap.String("docID", doc.ID), - zap.Error(err)) - return nil, common.CodeServerError, errors.New("Internal server error") - } - scheduledCount++ - } - - return map[string]interface{}{ - "scheduled_count": scheduledCount, - }, common.CodeSuccess, nil -} - -func (d *DatasetService) runEmbeddingDocument(kb *entity.Knowledgebase, doc *entity.Document, tableDoneCountByKB map[string]int64) error { - // Determine the flow ID: pipeline_id takes precedence, parser_id acts as - // fallback (builtin DSL from registry). Both paths go through the unified - // DSL ingestion worker. - flowID := "" - if doc.PipelineID != nil && strings.TrimSpace(*doc.PipelineID) != "" { - flowID = strings.TrimSpace(*doc.PipelineID) - } else if doc.ParserID != "" { - flowID = doc.ParserID - } - if flowID == "" { - return fmt.Errorf("document %s: no pipeline_id or parser_id configured", doc.ID) - } - - // Invalidate the cached column schema (field_map) when re-running table - // documents so the schema is rebuilt from the fresh parse result. - if doc.ParserID == string(entity.ParserTypeTable) { - doneCount, ok := tableDoneCountByKB[doc.KbID] - if !ok { - count, err := d.countDoneDocuments(doc.KbID) - if err != nil { - return err - } - doneCount = count - tableDoneCountByKB[doc.KbID] = doneCount - if doneCount <= 0 { - if err := d.kbDAO.DeleteFieldMap(doc.KbID); err != nil && !dao.IsNotFoundErr(err) { - return err - } - } - } - } - - return d.queueDatasetDataflowTask(kb, doc, flowID, 0) -} - -func (d *DatasetService) queueDatasetDataflowTask(kb *entity.Knowledgebase, doc *entity.Document, flowID string, priority int64) error { - if _, err := d.taskDAO.DeleteByDocIDs([]string{doc.ID}); err != nil { - return err - } - if err := d.beginDatasetParseDocument(doc.ID); err != nil { - return err - } - - now := time.Now() - task := &entity.Task{ - ID: utility.GenerateUUID(), - DocID: doc.ID, - FromPage: 0, - ToPage: maximumTaskPageNumber, - TaskType: "dataflow", - Priority: priority, - BeginAt: &now, - Progress: 0, - } - if err := d.taskDAO.CreateMany([]*entity.Task{task}); err != nil { - return err - } - - message := datasetParseTaskMessage(task) - message["task_type"] = task.TaskType - message["kb_id"] = doc.KbID - message["tenant_id"] = kb.TenantID - message["dataflow_id"] = flowID - message["file"] = nil - if redisClient := redisengine.Get(); redisClient == nil || !redisClient.QueueProduct(datasetParseQueueName(doc, priority), message) { - return fmt.Errorf("Can't access Redis. Please check the Redis' status.") - } - return nil -} - -func (d *DatasetService) countDoneDocuments(datasetID string) (int64, error) { - var count int64 - err := dao.GetDB().Model(&entity.Document{}). - Where("kb_id = ? AND run = ?", datasetID, string(entity.TaskStatusDone)). - Count(&count).Error - return count, err -} - -func (d *DatasetService) beginDatasetParseDocument(docID string) error { - now := time.Now() - return dao.GetDB().Model(&entity.Document{}).Where("id = ?", docID).Updates(map[string]interface{}{ - "progress_msg": "Task is queued...", - "process_begin_at": now, - "progress": rand.Float64() * 0.01, - "run": string(entity.TaskStatusRunning), - "chunk_num": 0, - "token_num": 0, - }).Error -} - // CheckEmbedding checks whether a new embedding model is compatible with stored vectors. func (d *DatasetService) CheckEmbedding(userID, datasetID string, req *CheckEmbeddingRequest) (*EmbeddingCheckResponse, common.ErrorCode, error) { if datasetID == "" { @@ -1190,35 +1052,6 @@ func datasetStringSlice(value interface{}) []string { } } -func datasetParseQueueName(doc *entity.Document, priority int64) string { - suffix := "common" - if doc.ParserID == string(entity.ParserTypeResume) { - suffix = "resume" - } - return fmt.Sprintf("%s.%d.%s", serverQueueNamePrefix, priority, suffix) -} - -func datasetParseTaskMessage(task *entity.Task) map[string]interface{} { - beginAt := "" - if task.BeginAt != nil { - beginAt = task.BeginAt.Format("2006-01-02 15:04:05") - } - digest := "" - if task.Digest != nil { - digest = *task.Digest - } - return map[string]interface{}{ - "id": task.ID, - "doc_id": task.DocID, - "from_page": task.FromPage, - "to_page": task.ToPage, - "progress": task.Progress, - "priority": task.Priority, - "begin_at": beginAt, - "digest": digest, - } -} - func (d *DatasetService) DeleteIndex(userID, datasetID, indexType string, wipe bool) (common.ErrorCode, error) { if !checkType(indexType) { return common.CodeArgumentError, fmt.Errorf("Invalid index type '%s'", indexType) diff --git a/test/testcases/restful_api/test_datasets.py b/test/testcases/restful_api/test_datasets.py index cb746984d1..74a9842930 100644 --- a/test/testcases/restful_api/test_datasets.py +++ b/test/testcases/restful_api/test_datasets.py @@ -2712,20 +2712,10 @@ def test_dataset_index_run_with_document_creates_task(rest_client, create_docume @pytest.mark.p2 -def test_dataset_embedding_endpoints(rest_client, create_dataset): - dataset_id = create_dataset("dataset_embedding_endpoints") - - run_no_docs_res = rest_client.post(f"/datasets/{dataset_id}/embedding") - assert run_no_docs_res.status_code == 200 - run_no_docs_payload = run_no_docs_res.json() - assert run_no_docs_payload["code"] == 102, run_no_docs_payload +def test_dataset_embedding_check_missing_embd_id(rest_client, create_dataset): + dataset_id = create_dataset("dataset_embedding_check_missing_embd_id") missing_embd_id_res = rest_client.post(f"/datasets/{dataset_id}/embedding/check", json={}) assert missing_embd_id_res.status_code == 200 missing_embd_id_payload = missing_embd_id_res.json() assert missing_embd_id_payload["code"] != 0, missing_embd_id_payload - - invalid_dataset_res = rest_client.post("/datasets/invalid_id/embedding") - assert invalid_dataset_res.status_code == 200 - invalid_dataset_payload = invalid_dataset_res.json() - assert invalid_dataset_payload["code"] != 0, invalid_dataset_payload diff --git a/test/testcases/test_http_api/common.py b/test/testcases/test_http_api/common.py index 54dd485f0b..3eff15f5b9 100644 --- a/test/testcases/test_http_api/common.py +++ b/test/testcases/test_http_api/common.py @@ -483,12 +483,6 @@ def delete_index(auth, dataset_id, index_type, *, headers=HEADERS): return res.json() -def run_embedding(auth, dataset_id, payload=None, *, headers=HEADERS): - url = f"{HOST_ADDRESS}{DATASETS_API_URL}/{dataset_id}/embedding" - res = requests.post(url=url, headers=headers, auth=auth, json=payload) - return res.json() - - def list_tags(auth, dataset_id, *, headers=HEADERS): url = f"{HOST_ADDRESS}{DATASETS_API_URL}/{dataset_id}/tags" res = requests.get(url=url, headers=headers, auth=auth) diff --git a/test/testcases/test_http_api/test_dataset_management/test_embedding.py b/test/testcases/test_http_api/test_dataset_management/test_embedding.py deleted file mode 100644 index 6ee5593962..0000000000 --- a/test/testcases/test_http_api/test_dataset_management/test_embedding.py +++ /dev/null @@ -1,32 +0,0 @@ -# -# Copyright 2025 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. -# -import pytest -from common import run_embedding - - -@pytest.mark.usefixtures("clear_datasets") -class TestRunEmbedding: - @pytest.mark.p2 - def test_run_embedding_no_documents(self, HttpApiAuth, add_dataset_func): - dataset_id = add_dataset_func - res = run_embedding(HttpApiAuth, dataset_id) - assert res["code"] == 102, res - assert "No documents in Dataset" in res.get("message", ""), res - - @pytest.mark.p2 - def test_run_embedding_invalid_id(self, HttpApiAuth): - res = run_embedding(HttpApiAuth, "invalid_id") - assert res["code"] != 0, res