youngkingdom/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-05-03 16:57:48 +08:00
Code Issues Packages Projects Releases Wiki Activity

Feat: Support get aggregated parsing status to dataset via the API (#13481)

Browse Source 
### What problem does this PR solve?

Support getting aggregated parsing status to dataset via the API

Issue: #12810

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

Co-authored-by: heyang.why <heyang.why@alibaba-inc.com>
This commit is contained in:
Heyang Wang
2026-03-10 18:05:45 +08:00
committed by GitHub
parent 68a623154a
commit 08f83ff331
7 changed files with 654 additions and 309 deletions
Download Patch File Download Diff File Expand all files Collapse all files

136
api/apps/sdk/dataset.py
View File

@ -138,12 +138,7 @@ async def create(tenant_id):
parser_cfg["metadata"] = fields
parser_cfg["enable_metadata"] = auto_meta.get("enabled", True)
req["parser_config"] = parser_cfg
e, req = KnowledgebaseService.create_with_name(
name = req.pop("name", None),
tenant_id = tenant_id,
parser_id = req.pop("parser_id", None),
**req
)
e, req = KnowledgebaseService.create_with_name(name=req.pop("name", None), tenant_id=tenant_id, parser_id=req.pop("parser_id", None), **req)
if not e:
return req
@ -159,19 +154,19 @@ async def create(tenant_id):
if not ok:
return err
try:
if not KnowledgebaseService.save(**req):
return get_error_data_result()
ok, k = KnowledgebaseService.get_by_id(req["id"])
if not ok:
return get_error_data_result(message="Dataset created failed")
response_data = remap_dictionary_keys(k.to_dict())
return get_result(data=response_data)
if not KnowledgebaseService.save(**req):
return get_error_data_result()
ok, k = KnowledgebaseService.get_by_id(req["id"])
if not ok:
return get_error_data_result(message="Dataset created failed")
response_data = remap_dictionary_keys(k.to_dict())
return get_result(data=response_data)
except Exception as e:
logging.exception(e)
return get_error_data_result(message="Database operation failed")
@manager.route("/datasets", methods=["DELETE"]) # noqa: F821
@token_required
async def delete(tenant_id):
@ -227,8 +222,7 @@ async def delete(tenant_id):
continue
kb_id_instance_pairs.append((kb_id, kb))
if len(error_kb_ids) > 0:
return get_error_permission_result(
message=f"""User '{tenant_id}' lacks permission for datasets: '{", ".join(error_kb_ids)}'""")
return get_error_permission_result(message=f"""User '{tenant_id}' lacks permission for datasets: '{", ".join(error_kb_ids)}'""")
errors = []
success_count = 0
@ -245,12 +239,12 @@ async def delete(tenant_id):
]
)
File2DocumentService.delete_by_document_id(doc.id)
FileService.filter_delete(
[File.source_type == FileSource.KNOWLEDGEBASE, File.type == "folder", File.name == kb.name])
FileService.filter_delete([File.source_type == FileSource.KNOWLEDGEBASE, File.type == "folder", File.name == kb.name])
# Drop index for this dataset
try:
from rag.nlp import search
idxnm = search.index_name(kb.tenant_id)
settings.docStoreConn.delete_idx(idxnm, kb_id)
except Exception as e:
@ -352,8 +346,7 @@ async def update(tenant_id, dataset_id):
try:
kb = KnowledgebaseService.get_or_none(id=dataset_id, tenant_id=tenant_id)
if kb is None:
return get_error_permission_result(
message=f"User '{tenant_id}' lacks permission for dataset '{dataset_id}'")
return get_error_permission_result(message=f"User '{tenant_id}' lacks permission for dataset '{dataset_id}'")
# Map auto_metadata_config into parser_config if present
auto_meta = req.pop("auto_metadata_config", None)
@ -384,8 +377,7 @@ async def update(tenant_id, dataset_id):
del req["parser_config"]
if "name" in req and req["name"].lower() != kb.name.lower():
exists = KnowledgebaseService.get_or_none(name=req["name"], tenant_id=tenant_id,
status=StatusEnum.VALID.value)
exists = KnowledgebaseService.get_or_none(name=req["name"], tenant_id=tenant_id, status=StatusEnum.VALID.value)
if exists:
return get_error_data_result(message=f"Dataset name '{req['name']}' already exists")
@ -393,8 +385,7 @@ async def update(tenant_id, dataset_id):
if not req["embd_id"]:
req["embd_id"] = kb.embd_id
if kb.chunk_num != 0 and req["embd_id"] != kb.embd_id:
return get_error_data_result(
message=f"When chunk_num ({kb.chunk_num}) > 0, embedding_model must remain {kb.embd_id}")
return get_error_data_result(message=f"When chunk_num ({kb.chunk_num}) > 0, embedding_model must remain {kb.embd_id}")
ok, err = verify_embedding_availability(req["embd_id"], tenant_id)
if not ok:
return err
@ -404,12 +395,10 @@ async def update(tenant_id, dataset_id):
return get_error_argument_result(message="'pagerank' can only be set when doc_engine is elasticsearch")
if req["pagerank"] > 0:
settings.docStoreConn.update({"kb_id": kb.id}, {PAGERANK_FLD: req["pagerank"]},
search.index_name(kb.tenant_id), kb.id)
settings.docStoreConn.update({"kb_id": kb.id}, {PAGERANK_FLD: req["pagerank"]}, search.index_name(kb.tenant_id), kb.id)
else:
# Elasticsearch requires PAGERANK_FLD be non-zero!
settings.docStoreConn.update({"exists": PAGERANK_FLD}, {"remove": PAGERANK_FLD},
search.index_name(kb.tenant_id), kb.id)
settings.docStoreConn.update({"exists": PAGERANK_FLD}, {"remove": PAGERANK_FLD}, search.index_name(kb.tenant_id), kb.id)
if not KnowledgebaseService.update_by_id(kb.id, req):
return get_error_data_result(message="Update dataset error.(Database error)")
@ -470,6 +459,15 @@ def list_datasets(tenant_id):
required: false
default: true
description: Order in descending.
- in: query
name: include_parsing_status
type: boolean
required: false
default: false
description: |
Whether to include document parsing status counts in the response.
When true, each dataset object will include: unstart_count, running_count,
cancel_count, done_count, and fail_count.
- in: header
name: Authorization
type: string
@ -487,17 +485,18 @@ def list_datasets(tenant_id):
if err is not None:
return get_error_argument_result(err)
include_parsing_status = args.get("include_parsing_status", False)
try:
kb_id = request.args.get("id")
name = args.get("name")
# check whether user has permission for the dataset with specified id
if kb_id:
kbs = KnowledgebaseService.get_kb_by_id(kb_id, tenant_id)
if not kbs:
if not KnowledgebaseService.get_kb_by_id(kb_id, tenant_id):
return get_error_permission_result(message=f"User '{tenant_id}' lacks permission for dataset '{kb_id}'")
# check whether user has permission for the dataset with specified name
if name:
kbs = KnowledgebaseService.get_kb_by_name(name, tenant_id)
if not kbs:
if not KnowledgebaseService.get_kb_by_name(name, tenant_id):
return get_error_permission_result(message=f"User '{tenant_id}' lacks permission for dataset '{name}'")
tenants = TenantService.get_joined_tenants_by_user_id(tenant_id)
@ -512,9 +511,17 @@ def list_datasets(tenant_id):
name,
)
parsing_status_map = {}
if include_parsing_status and kbs:
kb_ids = [kb["id"] for kb in kbs]
parsing_status_map = DocumentService.get_parsing_status_by_kb_ids(kb_ids)
response_data_list = []
for kb in kbs:
response_data_list.append(remap_dictionary_keys(kb))
data = remap_dictionary_keys(kb)
if include_parsing_status:
data.update(parsing_status_map.get(kb["id"], {}))
response_data_list.append(data)
return get_result(data=response_data_list, total=total)
except OperationalError as e:
logging.exception(e)
@ -530,9 +537,7 @@ def get_auto_metadata(tenant_id, dataset_id):
try:
kb = KnowledgebaseService.get_or_none(id=dataset_id, tenant_id=tenant_id)
if kb is None:
return get_error_permission_result(
message=f"User '{tenant_id}' lacks permission for dataset '{dataset_id}'"
)
return get_error_permission_result(message=f"User '{tenant_id}' lacks permission for dataset '{dataset_id}'")
parser_cfg = kb.parser_config or {}
metadata = parser_cfg.get("metadata") or []
@ -570,9 +575,7 @@ async def update_auto_metadata(tenant_id, dataset_id):
try:
kb = KnowledgebaseService.get_or_none(id=dataset_id, tenant_id=tenant_id)
if kb is None:
return get_error_permission_result(
message=f"User '{tenant_id}' lacks permission for dataset '{dataset_id}'"
)
return get_error_permission_result(message=f"User '{tenant_id}' lacks permission for dataset '{dataset_id}'")
parser_cfg = kb.parser_config or {}
fields = []
@ -598,20 +601,13 @@ async def update_auto_metadata(tenant_id, dataset_id):
return get_error_data_result(message="Database operation failed")
@manager.route('/datasets/<dataset_id>/knowledge_graph', methods=['GET']) # noqa: F821
@manager.route("/datasets/<dataset_id>/knowledge_graph", methods=["GET"]) # noqa: F821
@token_required
async def knowledge_graph(tenant_id, dataset_id):
if not KnowledgebaseService.accessible(dataset_id, tenant_id):
return get_result(
data=False,
message='No authorization.',
code=RetCode.AUTHENTICATION_ERROR
)
return get_result(data=False, message="No authorization.", code=RetCode.AUTHENTICATION_ERROR)
_, kb = KnowledgebaseService.get_by_id(dataset_id)
req = {
"kb_id": [dataset_id],
"knowledge_graph_kwd": ["graph"]
}
req = {"kb_id": [dataset_id], "knowledge_graph_kwd": ["graph"]}
obj = {"graph": {}, "mind_map": {}}
if not settings.docStoreConn.index_exist(search.index_name(kb.tenant_id), dataset_id):
@ -633,39 +629,29 @@ async def knowledge_graph(tenant_id, dataset_id):
obj["graph"]["nodes"] = sorted(obj["graph"]["nodes"], key=lambda x: x.get("pagerank", 0), reverse=True)[:256]
if "edges" in obj["graph"]:
node_id_set = {o["id"] for o in obj["graph"]["nodes"]}
filtered_edges = [o for o in obj["graph"]["edges"] if
o["source"] != o["target"] and o["source"] in node_id_set and o["target"] in node_id_set]
filtered_edges = [o for o in obj["graph"]["edges"] if o["source"] != o["target"] and o["source"] in node_id_set and o["target"] in node_id_set]
obj["graph"]["edges"] = sorted(filtered_edges, key=lambda x: x.get("weight", 0), reverse=True)[:128]
return get_result(data=obj)
@manager.route('/datasets/<dataset_id>/knowledge_graph', methods=['DELETE']) # noqa: F821
@manager.route("/datasets/<dataset_id>/knowledge_graph", methods=["DELETE"]) # noqa: F821
@token_required
def delete_knowledge_graph(tenant_id, dataset_id):
if not KnowledgebaseService.accessible(dataset_id, tenant_id):
return get_result(
data=False,
message='No authorization.',
code=RetCode.AUTHENTICATION_ERROR
)
return get_result(data=False, message="No authorization.", code=RetCode.AUTHENTICATION_ERROR)
_, kb = KnowledgebaseService.get_by_id(dataset_id)
settings.docStoreConn.delete({"knowledge_graph_kwd": ["graph", "subgraph", "entity", "relation"]},
search.index_name(kb.tenant_id), dataset_id)
settings.docStoreConn.delete({"knowledge_graph_kwd": ["graph", "subgraph", "entity", "relation"]}, search.index_name(kb.tenant_id), dataset_id)
return get_result(data=True)
@manager.route("/datasets/<dataset_id>/run_graphrag", methods=["POST"]) # noqa: F821
@token_required
def run_graphrag(tenant_id,dataset_id):
def run_graphrag(tenant_id, dataset_id):
if not dataset_id:
return get_error_data_result(message='Lack of "Dataset ID"')
if not KnowledgebaseService.accessible(dataset_id, tenant_id):
return get_result(
data=False,
message='No authorization.',
code=RetCode.AUTHENTICATION_ERROR
)
return get_result(data=False, message="No authorization.", code=RetCode.AUTHENTICATION_ERROR)
ok, kb = KnowledgebaseService.get_by_id(dataset_id)
if not ok:
@ -707,15 +693,11 @@ def run_graphrag(tenant_id,dataset_id):
@manager.route("/datasets/<dataset_id>/trace_graphrag", methods=["GET"]) # noqa: F821
@token_required
def trace_graphrag(tenant_id,dataset_id):
def trace_graphrag(tenant_id, dataset_id):
if not dataset_id:
return get_error_data_result(message='Lack of "Dataset ID"')
if not KnowledgebaseService.accessible(dataset_id, tenant_id):
return get_result(
data=False,
message='No authorization.',
code=RetCode.AUTHENTICATION_ERROR
)
return get_result(data=False, message="No authorization.", code=RetCode.AUTHENTICATION_ERROR)
ok, kb = KnowledgebaseService.get_by_id(dataset_id)
if not ok:
@ -734,15 +716,11 @@ def trace_graphrag(tenant_id,dataset_id):
@manager.route("/datasets/<dataset_id>/run_raptor", methods=["POST"]) # noqa: F821
@token_required
def run_raptor(tenant_id,dataset_id):
def run_raptor(tenant_id, dataset_id):
if not dataset_id:
return get_error_data_result(message='Lack of "Dataset ID"')
if not KnowledgebaseService.accessible(dataset_id, tenant_id):
return get_result(
data=False,
message='No authorization.',
code=RetCode.AUTHENTICATION_ERROR
)
return get_result(data=False, message="No authorization.", code=RetCode.AUTHENTICATION_ERROR)
ok, kb = KnowledgebaseService.get_by_id(dataset_id)
if not ok:
@ -784,7 +762,7 @@ def run_raptor(tenant_id,dataset_id):
@manager.route("/datasets/<dataset_id>/trace_raptor", methods=["GET"]) # noqa: F821
@token_required
def trace_raptor(tenant_id,dataset_id):
def trace_raptor(tenant_id, dataset_id):
if not dataset_id:
return get_error_data_result(message='Lack of "Dataset ID"')