Remove unused /datasets/<id>/embedding endpoint (#16936)

This commit is contained in:
euvre
2026-07-15 17:42:41 +08:00
committed by GitHub
parent 77d35a3387
commit 70a49c947d
8 changed files with 2 additions and 302 deletions

View File

@@ -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/<dataset_id>/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/<dataset_id>/embedding/check", methods=["POST"]) # noqa: F821
@login_required
@add_tenant_id_to_kwargs

View File

@@ -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.

View File

@@ -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) {

View File

@@ -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)

View File

@@ -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)

View File

@@ -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

View File

@@ -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)

View File

@@ -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