Feat: RAPTOR handle cancel gracefully (#11074)

### What problem does this PR solve?

RAPTOR handle cancel gracefully.

### Type of change

- [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
Yongteng Lei
2025-11-06 17:18:03 +08:00
committed by GitHub
parent 3bd1fefe1f
commit 0cd8024c34
3 changed files with 41 additions and 12 deletions

View File

@ -659,7 +659,7 @@ async def run_raptor_for_kb(row, kb_parser_config, chat_mdl, embd_mdl, vector_si
raptor_config["threshold"],
)
original_length = len(chunks)
chunks = await raptor(chunks, kb_parser_config["raptor"]["random_seed"], callback)
chunks = await raptor(chunks, kb_parser_config["raptor"]["random_seed"], callback, row["id"])
doc = {
"doc_id": fake_doc_id,
"kb_id": [str(row["kb_id"])],
@ -814,6 +814,8 @@ async def do_handle_task(task):
callback=progress_callback,
doc_ids=task.get("doc_ids", []),
)
if fake_doc_ids := task.get("doc_ids", []):
task_doc_id = fake_doc_ids[0] # use the first document ID to represent this task for logging purposes
# Either using graphrag or Standard chunking methods
elif task_type == "graphrag":
ok, kb = KnowledgebaseService.get_by_id(task_dataset_id)