mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-15 20:27:20 +08:00
Remove unused /datasets/<id>/embedding endpoint (#16936)
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user