From e3e2d0a3eca3993300937da04feae8aef7043442 Mon Sep 17 00:00:00 2001 From: buua436 Date: Thu, 23 Jul 2026 14:30:43 +0800 Subject: [PATCH] fix: optimize knowledge compilation LLM behavior and PageIndex ordering (#17270) --- api/apps/restful_apis/chunk_api.py | 44 +++++++++++++++++++ rag/advanced_rag/knowlege_compile/_common.py | 32 +++++++++++++- .../knowlege_compile/dataset_nav.py | 5 ++- rag/advanced_rag/knowlege_compile/raptor.py | 8 +++- rag/advanced_rag/knowlege_compile/runner.py | 2 + .../knowlege_compile/structure.py | 27 ++++++++---- rag/advanced_rag/knowlege_compile/wiki.py | 28 ++++++++++-- 7 files changed, 129 insertions(+), 17 deletions(-) diff --git a/api/apps/restful_apis/chunk_api.py b/api/apps/restful_apis/chunk_api.py index 151ba19f66..3c63e170f7 100644 --- a/api/apps/restful_apis/chunk_api.py +++ b/api/apps/restful_apis/chunk_api.py @@ -797,6 +797,50 @@ async def get_document_structure_graph(tenant_id, dataset_id, document_id): if bucket_id not in ordered_ids: ordered_ids.append(bucket_id) + page_index_groups = [group for group in grouped.values() if _compilation_template_kind(group.get("kind")) in {"page_index", "pageindex"}] + if page_index_groups: + # PageIndex nodes should follow the document's chunk order. Use the + # current chunk query order directly. Do not derive the order + # from page/position metadata because some document formats do not + # populate those fields. + chunk_order: dict[str, int] = {} + try: + chunk_res = await settings.retriever.search( + { + "doc_ids": [document_id], + "page": 1, + "size": 10000, + "question": "", + "sort": True, + "must_not": {"exists": "compile_kwd"}, + }, + index_name, + [dataset_id], + emb_mdl=None, + highlight=True, + ) + for index, chunk_id in enumerate(chunk_res.ids or []): + chunk_id = str(chunk_id) + if chunk_id: + chunk_order[chunk_id] = index + except Exception: + logging.exception("structure graph: failed to load document chunk order for PageIndex ordering") + + for group in page_index_groups: + original_entities = list(group.get("entities") or []) + + def entity_position(entity: dict, fallback: int): + chunk_indexes = [chunk_order[chunk_id] for chunk_id in entity.get("source_chunk_ids") or [] if isinstance(chunk_id, str) and chunk_id in chunk_order] + return (min(chunk_indexes), fallback) if chunk_indexes else (float("inf"), fallback) + + group["entities"] = [ + entity + for _, entity in sorted( + enumerate(original_entities), + key=lambda item: entity_position(item[1], item[0]), + ) + ] + templates_out = [grouped[bid] for bid in ordered_ids if grouped[bid]["entities"] or grouped[bid]["relations"]] return get_result(data={"templates": templates_out}) diff --git a/rag/advanced_rag/knowlege_compile/_common.py b/rag/advanced_rag/knowlege_compile/_common.py index 1adb0d9a53..b180d5779e 100644 --- a/rag/advanced_rag/knowlege_compile/_common.py +++ b/rag/advanced_rag/knowlege_compile/_common.py @@ -48,6 +48,31 @@ from rag.nlp import rag_tokenizer from rag.prompts.generator import INPUT_UTILIZATION, gen_json, split_chunks +def knowledge_compile_gen_conf(chat_mdl, gen_conf: Optional[dict] = None) -> dict: + """Add model-specific reasoning controls for knowledge compilation only.""" + conf = dict(gen_conf or {}) + model_config = getattr(chat_mdl, "model_config", None) + if not isinstance(model_config, dict): + model_config = {} + model_name = str(model_config.get("llm_name") or getattr(chat_mdl, "llm_name", "")).lower() + + if "deepseek-v4" in model_name: + extra_body = conf.get("extra_body") + extra_body = dict(extra_body) if isinstance(extra_body, dict) else {} + extra_body["thinking"] = {"type": "disabled"} + conf["extra_body"] = extra_body + elif "qwen3" in model_name: + # chat_model.py maps this flag to the provider-specific request body. + conf["enable_thinking"] = False + else: + # LiteLLM maps this common control for providers that support it and + # drops it for providers that do not. Keep model-specific overrides + # above because some APIs use a different, explicit request field. + conf["reasoning_effort"] = "none" + + return conf + + # --------------------------------------------------------------------------- # ID minting # --------------------------------------------------------------------------- @@ -772,7 +797,12 @@ async def _resolve_ambiguous_pairs( try: res = await asyncio.wait_for( - gen_json(system_prompt, user_prompt, chat_mdl, gen_conf={"temperature": 0.0}), + gen_json( + system_prompt, + user_prompt, + chat_mdl, + gen_conf=knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.0}), + ), timeout=llm_timeout, ) except asyncio.TimeoutError: diff --git a/rag/advanced_rag/knowlege_compile/dataset_nav.py b/rag/advanced_rag/knowlege_compile/dataset_nav.py index 54b249fa9e..b2dade9e5d 100644 --- a/rag/advanced_rag/knowlege_compile/dataset_nav.py +++ b/rag/advanced_rag/knowlege_compile/dataset_nav.py @@ -38,6 +38,7 @@ from common.misc_utils import thread_pool_exec from rag.utils.redis_conn import RedisDistributedLock from ._common import encode as _encode +from ._common import knowledge_compile_gen_conf as _knowledge_compile_gen_conf # --------------------------------------------------------------------------- # Constants @@ -465,7 +466,7 @@ async def _llm_merge(chat_mdl, cluster_desc: str, doc_summary: str) -> str: "Return ONLY the merged text, no commentary." ) try: - resp = await gen_json("", prompt, chat_mdl, gen_conf={"temperature": 0.1}) + resp = await gen_json("", prompt, chat_mdl, gen_conf=_knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.1})) if isinstance(resp, dict): return str(resp.get("merged", resp.get("result", cluster_desc))) if isinstance(resp, str) and resp.strip(): @@ -484,7 +485,7 @@ async def _llm_create_summary(chat_mdl, doc_summaries: list[str]) -> str: texts = "\n---\n".join(doc_summaries) prompt = f"Summarize the common topic of the following document excerpts in 1-3 concise sentences:\n\n{texts}\n\nReturn ONLY the summary text, no commentary." try: - resp = await gen_json("", prompt, chat_mdl, gen_conf={"temperature": 0.1}) + resp = await gen_json("", prompt, chat_mdl, gen_conf=_knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.1})) if isinstance(resp, dict): return str(resp.get("summary", resp.get("result", doc_summaries[0]))) if isinstance(resp, str) and resp.strip(): diff --git a/rag/advanced_rag/knowlege_compile/raptor.py b/rag/advanced_rag/knowlege_compile/raptor.py index cda5779e15..9dff4a749c 100644 --- a/rag/advanced_rag/knowlege_compile/raptor.py +++ b/rag/advanced_rag/knowlege_compile/raptor.py @@ -42,6 +42,7 @@ from rag.utils.raptor_utils import ( SUPPORTED_CLUSTERING_METHODS, SUPPORTED_TREE_BUILDERS, ) +from ._common import knowledge_compile_gen_conf # Regularization added to GMM covariance diagonals; keeps components # from collapsing on singleton/near-identical reduced points. @@ -216,7 +217,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval: last_exc = None for attempt in range(3): try: - response = await self._llm_model.async_chat(system, history, gen_conf) + response = await self._llm_model.async_chat(system, history, knowledge_compile_gen_conf(self._llm_model, gen_conf)) response = re.sub(r"^.*", "", response, flags=re.DOTALL) if response.find("**ERROR**") >= 0: raise Exception(response) @@ -863,7 +864,10 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval: break logging.info( "RAPTOR small-N collapse: layer of %d node(s) [%d:%d] collapsed into %d summary; stopping at tree top", - end - start, start, end, produced, + end - start, + start, + end, + produced, ) layers.append((end, len(chunks))) if callback: diff --git a/rag/advanced_rag/knowlege_compile/runner.py b/rag/advanced_rag/knowlege_compile/runner.py index c2ba8bb127..65f16addfb 100644 --- a/rag/advanced_rag/knowlege_compile/runner.py +++ b/rag/advanced_rag/knowlege_compile/runner.py @@ -85,6 +85,8 @@ def resolve_template_ids_from_groups(group_ids, tenant_id: str) -> list[str]: ids directly (the ``rag.flow`` Compiler carries them as a component parameter rather than inside ``parser_config``). """ + if isinstance(group_ids, str): + group_ids = [group_ids] template_ids: list[str] = [] seen: set[str] = set() for group_id in group_ids or []: diff --git a/rag/advanced_rag/knowlege_compile/structure.py b/rag/advanced_rag/knowlege_compile/structure.py index 7571020263..c8171a4ea4 100644 --- a/rag/advanced_rag/knowlege_compile/structure.py +++ b/rag/advanced_rag/knowlege_compile/structure.py @@ -35,10 +35,15 @@ from ._common import ( tokenize_for_search as _tokenize_for_search, union_ordered as _union_ordered, run_chunked_pipeline as _run_chunked_pipeline, + knowledge_compile_gen_conf as _knowledge_compile_gen_conf, ) _STRUCT_TYPES = ("list", "set", "hypergraph") +_STRUCT_TYPE_ALIASES = { + "graph": "hypergraph", + "knowledge_graph": "hypergraph", +} _ES_DEDUP_KNN_CONCURRENCY = 8 _ES_DEDUP_LLM_CONCURRENCY = 16 @@ -110,6 +115,7 @@ class PooledChatModel: return getattr(self._chat_mdl, name) async def async_chat(self, system, history, gen_conf=None, **kwargs): + gen_conf = _knowledge_compile_gen_conf(self._chat_mdl, gen_conf) return await self._pool.call( lambda: self._chat_mdl.async_chat(system, history, gen_conf=gen_conf, **kwargs), priority=self._priority, @@ -160,10 +166,12 @@ def _struct_get(cfg: dict, *keys, default=None): def _struct_infer_type(parser_config: dict) -> str: explicit = _struct_get(parser_config, "compile_type") normalized_explicit = _struct_normalize_kind(explicit) + normalized_explicit = _STRUCT_TYPE_ALIASES.get(normalized_explicit, normalized_explicit) if normalized_explicit in _STRUCT_TYPES: return normalized_explicit kind = _struct_get(parser_config, "kind") normalized_kind = _struct_normalize_kind(kind) + normalized_kind = _STRUCT_TYPE_ALIASES.get(normalized_kind, normalized_kind) if normalized_kind: return normalized_kind output = _struct_get(parser_config, "output", default={}) or {} @@ -176,7 +184,9 @@ def _struct_supported_type(parser_config: dict, autotype: str) -> bool: if autotype in _STRUCT_TYPES: return True kind = _struct_get(parser_config, "kind") - return _struct_normalize_kind(kind) == autotype + normalized_kind = _struct_normalize_kind(kind) + normalized_kind = _STRUCT_TYPE_ALIASES.get(normalized_kind, normalized_kind) + return normalized_kind == autotype def _struct_render_fields(fields: list, language: str) -> Tuple[str, str]: @@ -344,7 +354,7 @@ async def _struct_extract_hypergraph(text: str, parser_config: dict, chat_mdl, l node_prompt, edge_prompt_template = _struct_hypergraph_prompts(parser_config, language) user_prompt = f"## Source Text:\n{text}\n\n## Output (JSON only):" - node_res = await gen_json(node_prompt, user_prompt, chat_mdl, gen_conf={"temperature": 0.1}) + node_res = await gen_json(node_prompt, user_prompt, chat_mdl, gen_conf=_knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.1})) nodes = _struct_unwrap_items(node_res) id_field = _struct_entity_id_field(parser_config) @@ -362,7 +372,7 @@ async def _struct_extract_hypergraph(text: str, parser_config: dict, chat_mdl, l return nodes, [] edge_prompt = edge_prompt_template.replace("{known_nodes}", known_str) - edge_res = await gen_json(edge_prompt, user_prompt, chat_mdl, gen_conf={"temperature": 0.1}) + edge_res = await gen_json(edge_prompt, user_prompt, chat_mdl, gen_conf=_knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.1})) edges = _struct_unwrap_items(edge_res) return nodes, edges @@ -734,6 +744,7 @@ async def compile_structure_from_text( chunks, chat_mdl, prompt_overhead_tokens=prompt_overhead, + batch_size_cap=1 if str(template_kind).strip().lower().replace("-", "_") in {"page_index", "pageindex"} else None, ) if not packed_batches: return [] @@ -929,7 +940,7 @@ async def _struct_merge_pair(existing: dict, incoming: dict, chat_mdl) -> dict | item_incoming=json.dumps(incoming_payload, ensure_ascii=False), ) system_prompt = MERGE_SYSTEM_PROMPT + "\n\n" + MERGE_DECISION_INSTRUCTION - res = await gen_json(system_prompt, user_prompt, chat_mdl, gen_conf={"temperature": 0.0}) + res = await gen_json(system_prompt, user_prompt, chat_mdl, gen_conf=_knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.0})) if not isinstance(res, dict): return None if not res.get("duplicated"): @@ -1170,7 +1181,7 @@ async def _struct_judge_es_group_batch(group_specs: list[dict], chat_mdl) -> dic user_prompt = ES_GROUP_DECISION_BATCH_PROMPT.format(groups=json.dumps(prompt_groups, ensure_ascii=False)) system_prompt = MERGE_SYSTEM_PROMPT + "\n\n" + ES_GROUP_DECISION_BATCH_PROMPT.split("Groups:", 1)[0] - res = await gen_json(system_prompt, user_prompt, chat_mdl, gen_conf={"temperature": 0.0}) + res = await gen_json(system_prompt, user_prompt, chat_mdl, gen_conf=_knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.0})) raw_groups = res.get("groups") if isinstance(res, dict) else None if not isinstance(raw_groups, list): return {spec["request_group_id"]: set() for spec in group_specs} @@ -1221,7 +1232,7 @@ async def _struct_merge_es_group_batch(group_specs: list[dict], chat_mdl) -> dic user_prompt = ES_GROUP_BATCH_MERGE_PROMPT.format(groups=json.dumps(prompt_groups, ensure_ascii=False)) system_prompt = MERGE_SYSTEM_PROMPT + "\n\n" + ES_GROUP_BATCH_MERGE_PROMPT.split("Groups:", 1)[0] - res = await gen_json(system_prompt, user_prompt, chat_mdl, gen_conf={"temperature": 0.0}) + res = await gen_json(system_prompt, user_prompt, chat_mdl, gen_conf=_knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.0})) raw_groups = res.get("groups") if isinstance(res, dict) else None if not isinstance(raw_groups, list): return {spec["old_id"]: (list(spec["incoming_docs"]), None) for spec in group_specs} @@ -1277,7 +1288,7 @@ async def _struct_merge_es_group(old_doc: dict, incoming_docs: list[dict], chat_ ensure_ascii=False, ), ) - res = await gen_json(system_prompt, user_prompt, chat_mdl, gen_conf={"temperature": 0.0}) + res = await gen_json(system_prompt, user_prompt, chat_mdl, gen_conf=_knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.0})) if not isinstance(res, dict): return list(incoming_docs), None indices = res.get("duplicate_indices") @@ -2307,7 +2318,7 @@ async def validate_and_correct_chain( "You correct extracted graph relations to satisfy a strict-chain constraint.", prompt, chat_mdl, - gen_conf={"temperature": 0.0}, + gen_conf=_knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.0}), ) keep_raw = res.get("keep") if isinstance(res, dict) else None if isinstance(keep_raw, list): diff --git a/rag/advanced_rag/knowlege_compile/wiki.py b/rag/advanced_rag/knowlege_compile/wiki.py index 86ede89057..99d88bcd8b 100644 --- a/rag/advanced_rag/knowlege_compile/wiki.py +++ b/rag/advanced_rag/knowlege_compile/wiki.py @@ -49,6 +49,7 @@ from ._common import ( build_chunk_batches as _build_chunk_batches, bulk_dedup_items as _bulk_dedup_items, ensure_llm_bundle as _ensure_llm_bundle, + knowledge_compile_gen_conf as _knowledge_compile_gen_conf, run_chunked_pipeline as _run_chunked_pipeline, stable_row_id as _stable_row_id, ) @@ -839,7 +840,12 @@ async def _wiki_extract_one_batch( ) try: res = await asyncio.wait_for( - gen_json(WIKI_MAP_SYSTEM, user_prompt, chat_mdl, gen_conf={"temperature": 0.1}), + gen_json( + WIKI_MAP_SYSTEM, + user_prompt, + chat_mdl, + gen_conf=_knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.1}), + ), timeout=llm_timeout, ) except asyncio.TimeoutError: @@ -1930,7 +1936,12 @@ async def _wiki_resolve_maybe_items( try: res = await asyncio.wait_for( - gen_json(WIKI_PLAN_RECONCILE_SYSTEM, user_prompt, chat_mdl, gen_conf={"temperature": 0.0}), + gen_json( + WIKI_PLAN_RECONCILE_SYSTEM, + user_prompt, + chat_mdl, + gen_conf=_knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.0}), + ), timeout=llm_timeout, ) except asyncio.TimeoutError: @@ -2011,7 +2022,12 @@ async def _wiki_planning_call( try: res = await asyncio.wait_for( - gen_json(WIKI_PLAN_PLANNING_SYSTEM, user_prompt, chat_mdl, gen_conf={"temperature": 0.1}), + gen_json( + WIKI_PLAN_PLANNING_SYSTEM, + user_prompt, + chat_mdl, + gen_conf=_knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.1}), + ), timeout=llm_timeout, ) except asyncio.TimeoutError: @@ -3056,7 +3072,11 @@ async def _wiki_chat_text( logging.exception("wiki_refine: message_fit_in failed; sending untrimmed") try: raw = await asyncio.wait_for( - chat_mdl.async_chat(msg[0]["content"], msg[1:], {"temperature": temperature}), + chat_mdl.async_chat( + msg[0]["content"], + msg[1:], + _knowledge_compile_gen_conf(chat_mdl, {"temperature": temperature}), + ), timeout=llm_timeout, ) except asyncio.TimeoutError: