fix: optimize knowledge compilation LLM behavior and PageIndex ordering (#17270)

This commit is contained in:
buua436
2026-07-23 14:30:43 +08:00
committed by GitHub
parent c48a59db96
commit e3e2d0a3ec
7 changed files with 129 additions and 17 deletions

View File

@@ -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})

View File

@@ -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:

View File

@@ -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():

View File

@@ -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"^.*</think>", "", 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:

View File

@@ -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 []:

View File

@@ -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):

View File

@@ -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: