zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-14 00:48:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

Feat: Add knowledge compilation workflows (#16515)

Browse Source 
## Summary
- Add knowledge compilation template APIs, services, and builtin
template seed data
- Add advanced knowledge compile structure/artifact/RAPTOR workflow
support
- Update parsing, dataset/document APIs, and supporting services for
compilation workflows
This commit is contained in:
Kevin Hu
2026-07-02 23:22:07 +08:00
committed by GitHub
parent 7d64a78f83
commit 62f94cd59b
57 changed files with 14587 additions and 3094 deletions
Download Patch File Download Diff File Expand all files Collapse all files

517
rag/svr/task_executor_refactor/task_handler.py
View File

@@ -23,20 +23,26 @@ for handling document processing tasks with refactored, testable methods.
import asyncio
import logging
import json
# Wiki / artifact compilation pipeline lives in
# ``rag.svr.task_executor_refactor.dataset_wiki_generator`` — see the
# ``task_type == "artifact"`` branch of ``TaskHandler.run`` for the
# dispatch call.
# Document-structure compilation helpers (CHAIN_KINDS,
# compile_structure_from_text, merge_compiled_structures,
# validate_and_correct_chain) moved to ``chunk_post_processor``.
import xxhash
from timeit import default_timer as timer
from typing import Callable, Dict, List, Optional
from typing import AsyncIterator, Callable, Dict, List, Optional
from api.db.services.document_service import DocumentService
from api.db.services.knowledgebase_service import KnowledgebaseService
from api.db.services.compilation_template_group_service import CompilationTemplateGroupService
from api.db.joint_services.memory_message_service import handle_save_to_memory_task
from api.db.joint_services.tenant_model_service import (
get_tenant_default_model_by_type,
get_model_config_from_provider_instance
)
from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_from_provider_instance
from api.db.services.llm_service import LLMBundle
from api.db.services.task_service import GRAPH_RAPTOR_FAKE_DOC_ID
from api.db.services.task_service import GRAPH_RAPTOR_FAKE_DOC_ID, abort_doc_chunking_counter
from common.constants import LLMType
from common.exceptions import TaskCanceledException
from common.connection_utils import timeout
@@ -57,6 +63,96 @@ from rag.prompts.generator import run_toc_from_text
from common import settings
def _parser_config_compilation_template_group_ids(parser_config) -> list[str]:
"""Read template-group ids from a doc's parser_config.
Templates were previously referenced as a list
(``compilation_template_ids``); after the template-group refactor
a doc instead points at one or more groups, and the orchestrator
resolves each group's child templates at runtime. Old
``compilation_template_ids`` data is intentionally ignored per
the migration spec.
"""
def _normalize(raw) -> list[str]:
if isinstance(raw, str):
raw = [raw]
if not isinstance(raw, list):
return []
ids: list[str] = []
seen: set[str] = set()
for gid in raw:
if not isinstance(gid, str):
continue
gid = gid.strip()
if gid and gid not in seen:
seen.add(gid)
ids.append(gid)
return ids
if not isinstance(parser_config, dict):
return []
if "compilation_template_group_id" in parser_config:
return _normalize(parser_config.get("compilation_template_group_id"))
ext = parser_config.get("ext")
if isinstance(ext, dict):
return _normalize(ext.get("compilation_template_group_id"))
return []
def _parser_config_compilation_template_ids(parser_config, tenant_id: str) -> list[str]:
"""Resolve a doc's parser_config to compile-template ids by
looking up configured groups. Returns ``[]`` if the doc has no
group set or no group can be resolved.
"""
template_ids: list[str] = []
seen: set[str] = set()
for group_id in _parser_config_compilation_template_group_ids(parser_config):
for template_id in CompilationTemplateGroupService.resolve_template_ids(group_id, tenant_id):
if template_id in seen:
continue
seen.add(template_id)
template_ids.append(template_id)
return template_ids
def _resolve_template_chat_llm_id(parser_cfg: dict, ctx) -> str:
"""Pick the chat model id for a knowledge-compilation template.
Resolution order:
1. The template's own ``llm_id`` (what the user picked in the
compilation-template panel).
2. The doc's ``parser_config.llm_id`` (the doc-level chunking
model).
3. ``ctx.llm_id`` (the chunking task's default).
"""
if isinstance(parser_cfg, dict):
tid = parser_cfg.get("llm_id")
if isinstance(tid, str) and tid.strip():
return tid.strip()
doc_cfg = getattr(ctx, "parser_config", None) or {}
if isinstance(doc_cfg, dict):
did = doc_cfg.get("llm_id")
if isinstance(did, str) and did.strip():
return did.strip()
return ctx.llm_id
# Document-structure compilation tunables
# (DOC_STRUCTURE_COMPILE_BATCH_CHUNKS, DOC_STRUCTURE_MERGE_MAX_DOCS,
# STRUCTURE_CHAIN_CORRECTION_TIMEOUT_S) moved to
# ``chunk_post_processor``.
# Wiki / artifact tunables (``WIKI_MAP_BATCH_CHUNKS``,
# ``WIKI_GRAPH_MAX_CHUNK_IDS_PER_NODE``, commit-title / comments
# templates) moved to ``dataset_wiki_generator``.
# The corpus → skill compilation pipeline lives in
# ``rag.svr.task_executor_refactor.dataset_skill_generator``. Its entry
# point is :func:`run_corpus2skill`; this handler invokes it from the
# ``task_type == "skill"`` branch of ``run`` below.
class TaskHandler:
"""Main task handler for document processing.
@@ -85,32 +181,52 @@ class TaskHandler:
self._task_context = ctx
self._billing_hook = billing_hook
@staticmethod
def _is_standard_chunking_task(task_type: str) -> bool:
task_type = (task_type or "").lower()
return task_type not in {
"memory",
"raptor",
"graphrag",
"mindmap",
"artifact",
"skill",
"evaluation",
"reembedding",
"clone",
} and not task_type.startswith("dataflow")
async def handle_task(self) -> None:
try:
await self.handle()
except Exception:
if self._is_standard_chunking_task(self._task_context.task_type):
abort_doc_chunking_counter(self._task_context.doc_id)
raise
finally:
task_id = self._task_context.id
task_tenant_id = self._task_context.tenant_id
task_dataset_id = self._task_context.kb_id
task_doc_id = self._task_context.doc_id
if self._task_context.has_canceled_func(task_id):
try:
exists = await thread_pool_exec(
settings.docStoreConn.index_exist,
search.index_name(task_tenant_id),
task_dataset_id,
)
if exists:
ret = await thread_pool_exec(
settings.docStoreConn.delete,
{"doc_id": task_doc_id},
if self._is_standard_chunking_task(self._task_context.task_type):
abort_doc_chunking_counter(task_doc_id)
try:
exists = await thread_pool_exec(
settings.docStoreConn.index_exist,
search.index_name(task_tenant_id),
task_dataset_id,
)
self._task_context.recording_context.save_func_return_value("docStoreConn.delete", ret)
except Exception as e:
logging.exception(
f"Remove doc({task_doc_id}) from docStore failed when task({task_id}) canceled, exception: {e}")
if exists:
ret = await thread_pool_exec(
settings.docStoreConn.delete,
{"doc_id": task_doc_id},
search.index_name(task_tenant_id),
task_dataset_id,
)
self._task_context.recording_context.save_func_return_value("docStoreConn.delete", ret)
except Exception as e:
logging.exception(f"Remove doc({task_doc_id}) from docStore failed when task({task_id}) canceled, exception: {e}")
@timeout(60 * 60 * 3, 1)
async def handle(self) -> None:
@@ -134,7 +250,7 @@ class TaskHandler:
ctx.progress_cb(-1, msg="Task has been canceled.")
return
# Language defaults to "Chinese" via TaskContext._DEFAULTS — safe to bind model directly.
# Language defaults to "Chinese" via TaskContext._DEFAULTS 鈥?safe to bind model directly.
# Bind embedding model (matching original do_handle_task order: bind + init_kb before routing)
result = await self._bind_embedding_model()
if result is None:
@@ -154,12 +270,30 @@ class TaskHandler:
return
# Route to appropriate handler
if task_type == "raptor":
await self._run_raptor(embedding_model, vector_size)
elif task_type == "graphrag":
if task_type == "graphrag":
await self._run_graphrag(embedding_model)
elif task_type == "mindmap":
ctx.progress_cb(1, "place holder")
elif task_type == "artifact":
from rag.svr.task_executor_refactor.dataset_wiki_generator import (
run_wiki,
)
await run_wiki(
self._task_context,
embedding_model,
self._load_chunks_for_doc,
)
elif task_type == "skill":
from rag.svr.task_executor_refactor.dataset_skill_generator import (
run_corpus2skill,
)
await run_corpus2skill(
self._task_context,
embedding_model,
self._load_chunks_for_doc,
)
elif task_type == "evaluation":
await self._run_evaluation()
elif task_type == "reembedding":
@@ -167,8 +301,7 @@ class TaskHandler:
elif task_type == "clone":
await self._run_clone()
else:
await self._run_standard_chunking(embedding_model)
await self._run_standard_chunking(embedding_model, vector_size)
def _init_kb(self, vector_size: int) -> None:
"""Initialize knowledge base index."""
@@ -214,18 +347,14 @@ class TaskHandler:
try:
if task_embedding_id:
embd_model_config = get_model_config_from_provider_instance(
task_tenant_id, LLMType.EMBEDDING, task_embedding_id
)
embd_model_config = get_model_config_from_provider_instance(task_tenant_id, LLMType.EMBEDDING, task_embedding_id)
else:
embd_model_config = get_tenant_default_model_by_type(
task_tenant_id, LLMType.EMBEDDING
)
embd_model_config = get_tenant_default_model_by_type(task_tenant_id, LLMType.EMBEDDING)
embedding_model = LLMBundle(task_tenant_id, embd_model_config, lang=task_language)
vts, _ = embedding_model.encode(["ok"])
return embedding_model, len(vts[0])
except Exception as e:
error_message = f'Fail to bind embedding model: {str(e)}'
error_message = f"Fail to bind embedding model: {str(e)}"
ctx.progress_cb(-1, msg=error_message)
logging.exception(error_message)
raise
@@ -234,6 +363,7 @@ class TaskHandler:
self,
embedding_model: LLMBundle,
vector_size: int,
mark_done: bool = True,
) -> None:
"""Run RAPTOR summary generation."""
ctx = self._task_context
@@ -248,19 +378,21 @@ class TaskHandler:
kb_parser_config = kb.parser_config
if not kb_parser_config.get("raptor", {}).get("use_raptor", False):
kb_parser_config.update({
"raptor": {
"use_raptor": True,
"prompt": "Please summarize the following paragraphs. Be careful with the numbers, do not make things up. Paragraphs as following:\n {cluster_content}\nThe above is the content you need to summarize.",
"max_token": 256,
"threshold": 0.1,
"max_cluster": 64,
"random_seed": 0,
"scope": "file",
"clustering_method": "gmm",
"tree_builder": "raptor",
},
})
kb_parser_config.update(
{
"raptor": {
"use_raptor": True,
"prompt": "Please summarize the following paragraphs. Be careful with the numbers, do not make things up. Paragraphs as following:\n {cluster_content}\nThe above is the content you need to summarize.",
"max_token": 256,
"threshold": 0.1,
"max_cluster": 64,
"random_seed": 0,
"scope": "file",
"clustering_method": "gmm",
"tree_builder": "raptor",
},
}
)
if ctx.write_interceptor:
update_result = ctx.write_interceptor.intercept("KnowledgebaseService.update_by_id")
else:
@@ -271,11 +403,8 @@ class TaskHandler:
return
# Bind LLM for raptor
chat_model_config = get_model_config_from_provider_instance(
task_tenant_id, LLMType.CHAT, kb_task_llm_id
)
chat_model_config = get_model_config_from_provider_instance(task_tenant_id, LLMType.CHAT, kb_task_llm_id)
with LLMBundle(task_tenant_id, chat_model_config, lang=ctx.language) as chat_model:
# Run RAPTOR
raptor_service = RaptorService(ctx=ctx)
@@ -285,7 +414,7 @@ class TaskHandler:
chat_mdl=chat_model,
embd_mdl=embedding_model,
vector_size=vector_size,
doc_ids=ctx.doc_ids,
doc_ids=ctx.doc_ids or [ctx.doc_id],
)
ctx.recording_context.record("raptor_chunks", chunks)
@@ -293,7 +422,7 @@ class TaskHandler:
# Insert RAPTOR chunks
if chunks:
task_doc_id = (ctx.doc_ids or [GRAPH_RAPTOR_FAKE_DOC_ID])[0]
task_doc_id = (ctx.doc_ids or [ctx.doc_id] or [GRAPH_RAPTOR_FAKE_DOC_ID])[0]
chunk_service = ChunkService(ctx=ctx)
insert_result = await chunk_service.insert_chunks(ctx.id, task_tenant_id, task_dataset_id, chunks)
if insert_result:
@@ -304,27 +433,43 @@ class TaskHandler:
# Cleanup stale RAPTOR chunks
cleaned_chunks = 0
for cleanup_doc_id, keep_method in raptor_cleanup_chunks:
ret = await self._delete_raptor_chunks(
cleanup_doc_id, task_tenant_id, task_dataset_id, keep_method
)
ret = await self._delete_raptor_chunks(cleanup_doc_id, task_tenant_id, task_dataset_id, keep_method)
cleaned_chunks += ret
if cleaned_chunks:
ctx.progress_cb(msg=f"Cleaned up {cleaned_chunks} stale RAPTOR chunks.")
# Build the per-doc RAPTOR tree graph from the just-
# inserted summaries. Each chunk in ``chunks`` carries
# the doc_id it was written under (real doc id for
# scope="file"; GRAPH_RAPTOR_FAKE_DOC_ID for the
# dataset-scope path). We materialize one graph row per
# distinct doc_id so the dataset structure-graph
# endpoint can surface a RAPTOR tab per document.
# Failure here is best-effort — the summaries are
# already persisted; the tab just won't render.
raptor_doc_ids = {str(c.get("doc_id")) for c in chunks if c.get("doc_id")}
for raptor_doc_id in raptor_doc_ids:
try:
await raptor_service._persist_raptor_graph_to_es(raptor_doc_id)
except Exception:
logging.exception(
"raptor_graph: build failed for kb=%s doc=%s",
task_dataset_id,
raptor_doc_id,
)
# Update document stats
if ctx.write_interceptor:
ctx.write_interceptor.intercept("DocumentService.increment_chunk_num")
else:
DocumentService.increment_chunk_num(task_doc_id, task_dataset_id, token_count, len(chunks), 0)
ctx.recording_context.record("task_status", "completed")
ctx.progress_cb(prog=1.0, msg="RAPTOR done")
if mark_done:
ctx.recording_context.record("task_status", "completed")
ctx.progress_cb(prog=1.0, msg="RAPTOR done")
async def _run_graphrag(
self,
embedding_model: LLMBundle
) -> None:
async def _run_graphrag(self, embedding_model: LLMBundle) -> None:
"""Run GraphRAG."""
ctx = self._task_context
task_tenant_id = ctx.tenant_id
@@ -339,29 +484,31 @@ class TaskHandler:
kb_parser_config = kb.parser_config
if not kb_parser_config.get("graphrag", {}).get("use_graphrag", False):
kb_parser_config.update({
"graphrag": {
"use_graphrag": True,
"entity_types": [
"organization",
"person",
"geo",
"event",
"category",
],
"method": "light",
"batch_chunk_token_size": 4096,
"retry_attempts": 2,
"retry_backoff_seconds": 2.0,
"retry_backoff_max_seconds": 60.0,
"build_subgraph_timeout_per_chunk_seconds": 300,
"build_subgraph_min_timeout_seconds": 600,
"merge_timeout_seconds": 180,
"resolution_timeout_seconds": 1800,
"community_timeout_seconds": 1800,
"lock_acquire_timeout_seconds": 600,
kb_parser_config.update(
{
"graphrag": {
"use_graphrag": True,
"entity_types": [
"organization",
"person",
"geo",
"event",
"category",
],
"method": "light",
"batch_chunk_token_size": 4096,
"retry_attempts": 2,
"retry_backoff_seconds": 2.0,
"retry_backoff_max_seconds": 60.0,
"build_subgraph_timeout_per_chunk_seconds": 300,
"build_subgraph_min_timeout_seconds": 600,
"merge_timeout_seconds": 180,
"resolution_timeout_seconds": 1800,
"community_timeout_seconds": 1800,
"lock_acquire_timeout_seconds": 600,
}
}
})
)
if ctx.write_interceptor:
update_result = ctx.write_interceptor.intercept("KnowledgebaseService.update_by_id")
else:
@@ -372,11 +519,8 @@ class TaskHandler:
graphrag_conf = kb_parser_config.get("graphrag", {})
start_ts = timer()
chat_model_config = get_model_config_from_provider_instance(
task_tenant_id, LLMType.CHAT, kb_task_llm_id
)
chat_model_config = get_model_config_from_provider_instance(task_tenant_id, LLMType.CHAT, kb_task_llm_id)
with LLMBundle(task_tenant_id, chat_model_config, lang=task_language) as chat_model:
with_resolution = graphrag_conf.get("resolution", False)
with_community = graphrag_conf.get("community", False)
@@ -399,7 +543,20 @@ class TaskHandler:
async def _run_standard_chunking(
self,
embedding_model: LLMBundle
embedding_model: LLMBundle,
vector_size: int,
) -> None:
ctx = self._task_context
try:
await self._run_standard_chunking_impl(embedding_model, vector_size)
except Exception:
abort_doc_chunking_counter(ctx.doc_id)
raise
async def _run_standard_chunking_impl(
self,
embedding_model: LLMBundle,
vector_size: int,
) -> None:
"""Run standard chunking pipeline."""
ctx = self._task_context
@@ -409,7 +566,7 @@ class TaskHandler:
task_doc_id = ctx.doc_id
task_start_ts = timer()
doc_task_llm_id = ctx.parser_config.get("llm_id") or ctx.llm_id
ctx.raw_task['llm_id'] = doc_task_llm_id
ctx.raw_task["llm_id"] = doc_task_llm_id
# Build chunks
start_ts = timer()
@@ -419,9 +576,7 @@ class TaskHandler:
bucket, name = File2DocumentService.get_storage_address(doc_id=ctx.doc_id)
binary = await self._get_storage_binary(bucket, name)
if binary is None:
raise FileNotFoundError(
f"Can not find file <{ctx.name}> from minio. Could you try it again."
)
raise FileNotFoundError(f"Can not find file <{ctx.name}> from minio. Could you try it again.")
chunks = await chunk_service.build_chunks(binary)
ctx.recording_context.record("chunks", chunks)
@@ -431,7 +586,18 @@ class TaskHandler:
logging.info("Build document {}: {:.2f}s".format(ctx.name, timer() - start_ts))
if not chunks:
ctx.progress_cb(1., msg=f"No chunk built from {ctx.name}")
ctx.progress_cb(msg=f"No chunk built from {ctx.name}")
if not await self._run_document_post_chunking_if_last(
embedding_model,
vector_size,
task_start_ts,
0,
0,
):
return
task_time_cost = timer() - task_start_ts
ctx.recording_context.record("task_status", "completed")
ctx.progress_cb(prog=1.0, msg="Task done ({:.2f}s)".format(task_time_cost))
return
ctx.progress_cb(msg="Generate {} chunks".format(len(chunks)))
@@ -440,9 +606,7 @@ class TaskHandler:
start_ts = timer()
embedding_service = EmbeddingService(ctx=ctx)
try:
token_count, vector_size = await embedding_service.embed_chunks(
chunks, embedding_model, ctx.parser_config
)
token_count, vector_size = await embedding_service.embed_chunks(chunks, embedding_model, ctx.parser_config)
except TaskCanceledException:
raise
except Exception as e:
@@ -457,11 +621,6 @@ class TaskHandler:
logging.info(progress_message)
ctx.progress_cb(msg=progress_message)
# Build TOC if needed
toc_thread = None
if ctx.parser_id.lower() == "naive" and ctx.parser_config.get("toc_extraction", False):
toc_thread = asyncio.create_task(asyncio.to_thread(self._build_toc, ctx, chunks, ctx.progress_cb))
# Insert chunks
chunk_count = len(set([chunk["id"] for chunk in chunks]))
start_ts = timer()
@@ -469,15 +628,15 @@ class TaskHandler:
chunk_service = ChunkService(ctx=ctx)
if ctx.has_canceled_func(task_id):
abort_doc_chunking_counter(task_doc_id)
ctx.progress_cb(-1, msg="Task has been canceled.")
return
insert_result = await chunk_service.insert_chunks(
task_id, task_tenant_id, task_dataset_id, chunks
)
insert_result = await chunk_service.insert_chunks(task_id, task_tenant_id, task_dataset_id, chunks)
if not insert_result:
ctx.recording_context.record("insertion_result", "failed")
abort_doc_chunking_counter(task_doc_id)
return
ctx.recording_context.record("insertion_result", "success")
@@ -487,12 +646,8 @@ class TaskHandler:
ctx.progress_cb(msg="Indexing done ({:.2f}s).".format(timer() - start_ts))
toc_chunk = await self._process_toc_thread(toc_thread)
if toc_chunk:
ctx.recording_context.record("toc_chunk", [toc_chunk])
await post_processor.insert_toc_chunk(toc_chunk, chunk_service)
if ctx.has_canceled_func(task_id):
abort_doc_chunking_counter(task_doc_id)
ctx.progress_cb(-1, msg="Task has been canceled.")
return
@@ -502,17 +657,44 @@ class TaskHandler:
else:
DocumentService.increment_chunk_num(task_doc_id, task_dataset_id, token_count, chunk_count, 0)
if not await self._run_document_post_chunking_if_last(
embedding_model,
vector_size,
task_start_ts,
len(chunks),
token_count,
):
return
task_time_cost = timer() - task_start_ts
ctx.recording_context.record("task_status", "completed")
ctx.progress_cb(prog=1.0, msg="Task done ({:.2f}s)".format(task_time_cost))
logging.info(
"Chunk doc({}), page({}-{}), chunks({}), token({}), elapsed:{:.2f}".format(
ctx.name, ctx.from_page, ctx.to_page,
len(chunks), token_count, task_time_cost
)
logging.info("Chunk doc({}), page({}-{}), chunks({}), token({}), elapsed:{:.2f}".format(ctx.name, ctx.from_page, ctx.to_page, len(chunks), token_count, task_time_cost))
async def _run_document_post_chunking_if_last(
self,
embedding_model: LLMBundle,
vector_size: int,
task_start_ts: float,
chunks_len: int,
token_count: int,
) -> bool:
"""Thin delegator. The pipeline lives in
``rag.svr.task_executor_refactor.chunk_post_processor``.
"""
from rag.svr.task_executor_refactor.chunk_post_processor import (
run_document_post_chunking_if_last,
)
return await run_document_post_chunking_if_last(
self,
embedding_model,
vector_size,
task_start_ts,
chunks_len,
token_count,
)
async def _process_toc_thread(self, toc_thread):
try:
@@ -527,30 +709,105 @@ class TaskHandler:
@classmethod
async def _get_storage_binary(cls, bucket: str, name: str) -> bytes:
from common import settings
"""Get binary from storage."""
return await thread_pool_exec(settings.STORAGE_IMPL.get, bucket, name)
@staticmethod
async def _load_chunks_for_doc(
tenant_id: str,
kb_id: str,
doc_id: str,
batch_size: int = 500,
) -> AsyncIterator[List[Dict]]:
"""Stream a document's chunks from the doc store one batch at a time.
Async generator that yields successive batches of up to ``batch_size``
chunks. Order is pushed to the doc store via
``OrderByExpr().asc("page_num_int").asc("top_int")`` so callers do
not need to re-sort. Rows with a ``compile_kwd`` marker (artifact
pages, structure entities, etc.) are filtered out defensively.
Memory is bounded by ``batch_size``: at most one page is materialised
at a time, so long documents do not balloon the worker's heap.
"""
from common.doc_store.doc_store_base import OrderByExpr
index_nm = search.index_name(tenant_id)
if not settings.docStoreConn.index_exist(index_nm, kb_id):
return
select_fields = [
"id",
"doc_id",
"content_with_weight",
"page_num_int",
"top_int",
]
order_by = OrderByExpr()
order_by.asc("page_num_int")
order_by.asc("top_int")
offset = 0
while True:
try:
res = await thread_pool_exec(
settings.docStoreConn.search,
select_fields,
[],
{"doc_id": [doc_id], "available_int": 1},
[],
order_by,
offset,
batch_size,
index_nm,
[kb_id],
)
field_map = settings.docStoreConn.get_fields(res, select_fields)
except Exception:
logging.exception("load_chunks_for_doc: failed to load chunks for doc=%s", doc_id)
return
if not field_map:
return
batch: List[Dict] = []
for row_id, row in field_map.items():
if row.get("compile_kwd"):
continue
batch.append(
{
"id": row_id,
"doc_id": row.get("doc_id") or doc_id,
"content_with_weight": row.get("content_with_weight") or "",
"page_num_int": row.get("page_num_int", 0),
"top_int": row.get("top_int", 0),
}
)
if batch:
yield batch
if len(field_map) < batch_size:
return
offset += batch_size
@classmethod
def _build_toc(cls, ctx: TaskContext, docs: List[Dict], progress_cb: Callable) -> Optional[Dict]:
"""Build table of contents."""
progress_cb(msg="Start to generate table of content ...")
chat_model_config = get_model_config_from_provider_instance(
ctx.tenant_id, LLMType.CHAT, ctx.llm_id
)
chat_model_config = get_model_config_from_provider_instance(ctx.tenant_id, LLMType.CHAT, ctx.llm_id)
with LLMBundle(ctx.tenant_id, chat_model_config, lang=ctx.language) as chat_mdl:
docs = sorted(docs, key=lambda d: (
d.get("page_num_int", 0)[0] if isinstance(d.get("page_num_int", 0), list) else d.get("page_num_int", 0),
d.get("top_int", 0)[0] if isinstance(d.get("top_int", 0), list) else d.get("top_int", 0)
))
docs = sorted(
docs,
key=lambda d: (
d.get("page_num_int", 0)[0] if isinstance(d.get("page_num_int", 0), list) else d.get("page_num_int", 0),
d.get("top_int", 0)[0] if isinstance(d.get("top_int", 0), list) else d.get("top_int", 0),
),
)
# NOTE: asyncio.run() creates a new event loop in the worker thread
# (this method is called via asyncio.to_thread), which is the
# intended pattern for bridging sync -> async in a thread context.
toc: list[dict] = asyncio.run(
run_toc_from_text([d["content_with_weight"] for d in docs], chat_mdl, progress_cb)
)
logging.info("------------ T O C -------------\n" + json.dumps(toc, ensure_ascii=False, indent=' '))
toc: list[dict] = asyncio.run(run_toc_from_text([d["content_with_weight"] for d in docs], chat_mdl, progress_cb))
logging.info("------------ T O C -------------\n" + json.dumps(toc, ensure_ascii=False, indent=" "))
for ii, item in enumerate(toc):
try:
@@ -578,19 +835,17 @@ class TaskHandler:
if toc:
import copy
d = copy.deepcopy(docs[-1])
d["content_with_weight"] = json.dumps(toc, ensure_ascii=False)
d["toc_kwd"] = "toc"
d["available_int"] = 0
d["page_num_int"] = [100000000]
d["id"] = xxhash.xxh64(
(d["content_with_weight"] + str(d["doc_id"])).encode("utf-8", "surrogatepass")).hexdigest()
d["id"] = xxhash.xxh64((d["content_with_weight"] + str(d["doc_id"])).encode("utf-8", "surrogatepass")).hexdigest()
return d
return None
async def _delete_raptor_chunks(
self, doc_id: str, tenant_id: str, kb_id: str, keep_method: Optional[str]
) -> int:
async def _delete_raptor_chunks(self, doc_id: str, tenant_id: str, kb_id: str, keep_method: Optional[str]) -> int:
"""Delete RAPTOR chunks."""
if self._task_context.write_interceptor:
return self._task_context.write_interceptor.intercept("delete_raptor_chunks")