ragflow/rag/flow/compiler/compiler.py

#
#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
import logging
import random
from copy import deepcopy
from types import SimpleNamespace

import xxhash

from agent.component.llm import LLMParam, LLM
from api.db.joint_services.tenant_model_service import get_model_config_by_id, get_tenant_default_model_by_type, resolve_model_config
from api.db.services.document_service import DocumentService
from api.db.services.knowledgebase_service import KnowledgebaseService
from api.db.services.llm_service import LLMBundle
from api.db.services.task_service import has_canceled
from common.constants import LLMType
from rag.advanced_rag.knowlege_compile.runner import (
    DOC_STRUCTURE_COMPILE_BATCH_CHUNKS,
    load_active_templates,
    resolve_template_ids_from_groups,
    run_structure_compile_over_batches,
    split_tree_templates,
)
from rag.flow.base import ProcessBase, ProcessParamBase


class CompilerParam(ProcessParamBase, LLMParam):
    """Parameters for the knowledge-Compiler flow component.

    Same LLM-backed shape as the Extractor, but instead of a single inline
    ``knowledge_compilation`` config it drives compilation from one or more
    saved **compilation-template groups** (``compilation_template_group_ids``).
    Each group resolves to a set of templates, each of which carries its own
    structure-compilation config (kind, fields, synthesis, ...).
    """

    def __init__(self):
        super().__init__()
        self.compilation_template_group_ids = []

    def check(self):
        super().check()
        self.check_empty(self.compilation_template_group_ids, "Compilation Template Groups")


class Compiler(ProcessBase, LLM):
    component_name = "Compiler"

    def _compile_progress(self, prog=None, msg=""):
        """Adapt the knowledge-compile ``callback`` protocol to the flow
        callback. Downstream compile helpers invoke the callback either as
        ``callback(prog, msg)`` (positional) or ``callback(msg=...)``; the
        flow's ``self.callback`` expects ``(progress, message)``.
        """
        self.callback(0 if prog is None else prog, msg)

    async def _compile_tree_templates(
        self,
        templates: list[tuple[str, dict]],
        chat_mdl_by_tid: dict[str, LLMBundle],
        embedding_model: LLMBundle,
        chunks: list[dict],
        tenant_id: str,
        kb_id: str,
        doc_id: str,
    ) -> None:
        """Build and persist tree graphs from the pipeline's in-memory chunks.

        The document post-chunking path can reload chunks from the doc store,
        but a pipeline Compiler runs before DataflowService persists its final
        chunks. Supply RAPTOR with the same ``(text, vector, chunk_id)`` shape
        from the current pipeline output instead.
        """
        from rag.advanced_rag.knowlege_compile.structure import _struct_upsert_graph_json
        from rag.svr.task_executor_refactor.chunk_post_processor import raptor_tree_to_graph
        from rag.svr.task_executor_refactor.raptor_service import RaptorService

        tree_inputs = []
        texts = []
        for chunk in chunks:
            text = chunk.get("content_with_weight") or chunk.get("text") or ""
            if not isinstance(text, str) or not text.strip():
                continue
            chunk_id = str(chunk.get("id") or "")
            if not chunk_id:
                continue
            texts.append(text)
            tree_inputs.append((text, chunk_id))
        if not tree_inputs:
            return

        vectors, _ = embedding_model.encode(texts)
        tree_chunks = [(text, vector, chunk_id) for (text, chunk_id), vector in zip(tree_inputs, vectors)]
        if not tree_chunks:
            return

        tree_context = SimpleNamespace(
            tenant_id=tenant_id,
            kb_id=kb_id,
            doc_id=doc_id,
            id=getattr(self._canvas, "task_id", ""),
            progress_cb=self._compile_progress,
        )
        raptor_service = RaptorService(tree_context)

        for idx, (template_id, parser_cfg) in enumerate(templates):
            raptor_cfg = (parser_cfg or {}).get("raptor") or {}
            raptor_config = {
                "prompt": raptor_cfg.get("prompt") or "Please write a concise summary of the following texts:\n{cluster_content}",
                "max_token": int(raptor_cfg.get("max_token") or 512),
                "threshold": float(raptor_cfg.get("threshold") or 0.1),
                "random_seed": int(raptor_cfg.get("random_seed") or 0),
                "max_cluster": int(raptor_cfg.get("max_cluster") or 64),
                "ext": raptor_cfg.get("ext") or {},
            }
            self._compile_progress(msg=f"tree-template ({idx + 1}/{len(templates)}): building tree for doc={doc_id}")
            try:
                tree = await raptor_service.build_doc_tree(
                    chunks=tree_chunks,
                    raptor_config=raptor_config,
                    chat_mdl=chat_mdl_by_tid[template_id],
                    embd_mdl=embedding_model,
                    tree_builder="raptor",
                    clustering_method="gmm",
                    max_errors=3,
                )
            except Exception:
                logging.exception("Compiler: tree-template %s build failed for doc %s", template_id, doc_id)
                continue
            if tree is None:
                continue

            if bool(raptor_cfg.get("rechunk")):
                self._compile_progress(msg="Compiler: tree rechunking is not supported for in-memory pipeline chunks; keeping original chunks.")

            try:
                await _struct_upsert_graph_json(
                    raptor_tree_to_graph(tree),
                    tenant_id,
                    kb_id,
                    doc_id,
                    compile_kwd="tree",
                    compilation_template_id=template_id,
                )
            except Exception:
                logging.exception("Compiler: tree-template %s graph upsert failed for doc %s", template_id, doc_id)
                continue

            try:
                from rag.advanced_rag.knowlege_compile.dataset_nav import upsert_dataset_nav_doc

                await upsert_dataset_nav_doc(tenant_id, kb_id, doc_id, tree)
            except Exception:
                logging.exception("Compiler: tree-template %s dataset navigation upsert failed for doc %s", template_id, doc_id)

            self._compile_progress(msg=f"tree-template ({idx + 1}/{len(templates)}): persisted tree graph for doc {doc_id}")

    def _compile_language(self, kwargs: dict) -> str:
        language = kwargs.get("language") or getattr(self._canvas, "_language", None)
        if isinstance(language, str):
            language = language.strip()
        if not language and getattr(self._canvas, "_doc_id", None):
            config = DocumentService.get_chunking_config(self._canvas._doc_id) or {}
            language = config.get("language")
            if isinstance(language, str):
                language = language.strip()
        return language or "English"

    async def _invoke(self, **kwargs):
        self.set_output("output_format", "chunks")
        self.callback(random.randint(1, 5) / 100.0, "Start knowledge compilation.")

        # Pipeline components receive the previous component's output as
        # kwargs. Do not call LLM.get_input_elements() here: it resolves the
        # inherited prompt variables through Canvas.globals, while Pipeline
        # is a Graph and has no globals.
        chunks = deepcopy(kwargs.get("chunks") or [])
        if not chunks:
            for val in kwargs.values():
                if isinstance(val, list):
                    chunks = deepcopy(val)

        tenant_id = self._canvas.get_tenant_id()
        doc_id = self._canvas._doc_id
        kb_id = getattr(self._canvas, "_kb_id", None) or DocumentService.get_knowledgebase_id(doc_id)
        language = self._compile_language(kwargs)

        if not chunks:
            self.set_output("chunks", chunks)
            return

        for ck in chunks:
            ck["doc_id"] = doc_id
            ck["id"] = xxhash.xxh64((ck["text"] + str(ck["doc_id"])).encode("utf-8")).hexdigest()

        # Resolve the configured template groups to concrete, active
        # (non-artifact) structure-compilation templates.
        template_ids = resolve_template_ids_from_groups(self._param.compilation_template_group_ids, tenant_id)
        active_templates = load_active_templates(template_ids, tenant_id)
        if not active_templates:
            self.callback(0, "No active compilation templates resolved from the configured groups.")
            self.set_output("chunks", chunks)
            return

        # Per-template chat model: a template may pin its own ``llm_id``;
        # otherwise fall back to this component's configured chat model.
        llm_bundle_cache: dict[str, LLMBundle] = {}
        chat_mdl_by_tid: dict[str, LLMBundle] = {}
        filtered_templates: list[tuple[str, dict]] = []
        default_chat_mdl = None
        for template_id, parser_cfg in active_templates:
            tpl_llm_id = parser_cfg.get("llm_id") if isinstance(parser_cfg, dict) else None
            if isinstance(tpl_llm_id, str) and tpl_llm_id.strip():
                chat_llm_id = tpl_llm_id.strip()
                if chat_llm_id not in llm_bundle_cache:
                    try:
                        cfg = resolve_model_config(tenant_id, LLMType.CHAT, chat_llm_id)
                        llm_bundle_cache[chat_llm_id] = LLMBundle(
                            tenant_id,
                            cfg,
                            lang=language,
                            max_retries=self._param.max_retries,
                            retry_interval=self._param.delay_after_error,
                        )
                    except Exception:
                        logging.exception(
                            "Compiler: cannot resolve chat model %s for template %s; skipping",
                            chat_llm_id,
                            template_id,
                        )
                        continue
                chat_mdl_by_tid[template_id] = llm_bundle_cache[chat_llm_id]
            else:
                if default_chat_mdl is None:
                    default_chat_mdl = LLMBundle(
                        tenant_id,
                        self.chat_mdl.model_config,
                        lang=language,
                        max_retries=self._param.max_retries,
                        retry_interval=self._param.delay_after_error,
                    )
                chat_mdl_by_tid[template_id] = default_chat_mdl
            filtered_templates.append((template_id, parser_cfg))

        if not filtered_templates:
            self.set_output("chunks", chunks)
            return
        active_templates = filtered_templates

        if self._canvas._kb_id:
            e, kb = KnowledgebaseService.get_by_id(self._canvas._kb_id)
            if kb.tenant_embd_id:
                try:
                    embd_model_config = get_model_config_by_id(self._canvas._tenant_id, LLMType.EMBEDDING, kb.tenant_embd_id)
                except LookupError:
                    embd_model_config = resolve_model_config(self._canvas._tenant_id, LLMType.EMBEDDING, kb.embd_id)
            else:
                embd_model_config = resolve_model_config(self._canvas._tenant_id, LLMType.EMBEDDING, kb.embd_id)
        else:
            embd_model_config = get_tenant_default_model_by_type(self._canvas._tenant_id, LLMType.EMBEDDING)
        embedding_model = LLMBundle(
            tenant_id,
            embd_model_config,
            lang=language,
            max_retries=self._param.max_retries,
            retry_interval=self._param.delay_after_error,
        )

        tree_templates, non_tree_templates = split_tree_templates(active_templates)
        if tree_templates:
            await self._compile_tree_templates(
                tree_templates,
                chat_mdl_by_tid,
                embedding_model,
                chunks,
                tenant_id,
                kb_id,
                doc_id,
            )

        if non_tree_templates:
            task_id = getattr(self._canvas, "task_id", None)

            def _cancelled() -> bool:
                return bool(task_id) and has_canceled(task_id)

            async def _chunk_batches():
                for i in range(0, len(chunks), DOC_STRUCTURE_COMPILE_BATCH_CHUNKS):
                    yield chunks[i : i + DOC_STRUCTURE_COMPILE_BATCH_CHUNKS]

            await run_structure_compile_over_batches(
                active_templates=non_tree_templates,
                chat_mdl_by_tid=chat_mdl_by_tid,
                embedding_model=embedding_model,
                tenant_id=tenant_id,
                kb_id=kb_id,
                doc_id=doc_id,
                language=language,
                chunk_batches=_chunk_batches(),
                progress_cb=self._compile_progress,
                cancel_check=_cancelled,
            )

        self.set_output("chunks", chunks)