diff --git a/admin/client/parser.py b/admin/client/parser.py index 8f91352ad1..1cdb206b26 100644 --- a/admin/client/parser.py +++ b/admin/client/parser.py @@ -112,6 +112,8 @@ sql_command: login_user | set_license | set_license_config | show_license + | generate_nav_for_dataset + | navigation_search | check_license | benchmark @@ -171,6 +173,9 @@ ENVS: "ENVS"i KEY: "KEY"i KEYS: "KEYS"i GENERATE: "GENERATE"i +NAVIGATION: "NAVIGATION"i +MODE: "MODE"i +TOPK: "TOPK"i MODEL: "MODEL"i MODELS: "MODELS"i PROVIDER: "PROVIDER"i @@ -194,6 +199,7 @@ SIZE: "SIZE"i PARSER: "PARSER"i PIPELINE: "PIPELINE"i SEARCH: "SEARCH"i +EXPLORE: "EXPLORE"i CURRENT: "CURRENT"i LLM: "LLM"i VLM: "VLM"i @@ -373,6 +379,8 @@ create_metadata_table: CREATE METADATA TABLE ";" drop_metadata_table: DROP METADATA TABLE ";" insert_dataset_from_file: INSERT DATASET FROM FILE quoted_string ";" insert_metadata_from_file: INSERT METADATA FROM FILE quoted_string ";" +generate_nav_for_dataset: GENERATE NAVIGATION OF DATASET quoted_string ";" +navigation_search: NAVIGATION SEARCH quoted_string IN DATASET quoted_string MODE quoted_string (TOPK NUMBER)? ";" update_chunk: UPDATE CHUNK quoted_string OF DATASET quoted_string SET quoted_string ";" identifier_list: identifier (COMMA identifier)* @@ -788,6 +796,35 @@ class RAGFlowCLITransformer(Transformer): file_path = items[4].children[0].strip("'\"") return {"type": "insert_metadata_from_file", "file_path": file_path} + def generate_nav_for_dataset(self, items): + dataset_id = items[4].children[0].strip("'\"") + return {"type": "generate_nav_for_dataset", "dataset_id": dataset_id} + + def navigation_search(self, items): + query = items[2].children[0].strip("'\"") + dataset_name = items[5].children[0].strip("'\"") + mode = items[7].children[0].strip("'\"") + topk = None + # If TOPK was specified, items will be longer. + if len(items) > 9: + try: + extra = items[8] + if hasattr(extra, "data"): + # group wrapper: extra.children = [TOPK, NUMBER] + topk = int(extra.children[1]) + else: + # flattened: items[8] = TOPK, items[9] = NUMBER + topk = int(items[9]) + except (ValueError, IndexError): + topk = None + return { + "type": "navigation_search", + "query": query, + "dataset_id": dataset_name, + "mode": mode, + "topk": topk, + } + def update_chunk(self, items): def get_quoted_value(item): if hasattr(item, "children") and item.children: diff --git a/admin/client/ragflow_client.py b/admin/client/ragflow_client.py index 8b67537b1a..a30d64a6a8 100644 --- a/admin/client/ragflow_client.py +++ b/admin/client/ragflow_client.py @@ -13,21 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import base64 import json -import time -from typing import Any, List, Optional import multiprocessing as mp -from concurrent.futures import ProcessPoolExecutor, as_completed +import time import urllib.parse +from concurrent.futures import ProcessPoolExecutor, as_completed from pathlib import Path +from typing import Any + +from Cryptodome.Cipher import PKCS1_v1_5 as Cipher_pkcs1_v1_5 +from Cryptodome.PublicKey import RSA from http_client import HttpClient from lark import Tree from user import encrypt_password, login_user -import base64 -from Cryptodome.Cipher import PKCS1_v1_5 as Cipher_pkcs1_v1_5 -from Cryptodome.PublicKey import RSA - try: from requests_toolbelt import MultipartEncoder except Exception as e: # pragma: no cover - fallback without toolbelt @@ -129,7 +129,6 @@ class RAGFlowClient: print(f"Fail to get all services, code: {res_json['code']}, message: {res_json['message']}") else: print(f"Fail to get all services, code: {response.status_code}, body: {response.text}") - pass def show_service(self, command): if self.server_type != "admin": @@ -1802,6 +1801,76 @@ class RAGFlowClient: else: print(f"Fail to set metadata, code: {response.status_code}, body: {response.text}") + def generate_nav_for_dataset(self, command_dict): + if self.server_type != "user": + print("This command is only allowed in USER mode") + return + + dataset_name = command_dict["dataset_id"] + dataset_id = self._get_dataset_id(dataset_name) + if dataset_id is None: + print(f"Dataset not found: {dataset_name}") + return + + response = self.http_client.request( + "POST", + f"/datasets/{dataset_id}/navigation", + json_body={}, + use_api_base=True, + auth_kind="web", + ) + if response.status_code == 200: + res_json = response.json() + if res_json.get("code") == 0: + data = res_json.get("data", {}) + print(f"Navigation tree created: deleted={data.get('deleted', 0)}, upserted={data.get('upserted', 0)}") + else: + print(f"Fail to generate navigation, code: {res_json.get('code')}, message: {res_json.get('message')}") + else: + print(f"Fail to generate navigation, code: {response.status_code}, body: {response.text}") + + def navigation_search(self, command_dict): + if self.server_type != "user": + print("This command is only allowed in USER mode") + return + + query = command_dict["query"] + dataset_name = command_dict["dataset_id"] + mode = command_dict["mode"] + topk = command_dict.get("topk", None) + + # Try name lookup; fall back to using the value directly as an ID. + dataset_id = self._get_dataset_id(dataset_name) + if dataset_id is None: + dataset_id = dataset_name + + valid_modes = {"nav_doc", "nav_cluster", "navigation_tree", "chunk", "all"} + if mode not in valid_modes: + print(f"Invalid mode: {mode}, expected one of {valid_modes}") + return + + url = f"/datasets/{dataset_id}/navigation/search?q={query}&mode={mode}" + if topk is not None: + url += f"&top_k={topk}" + response = self.http_client.request( + "GET", + url, + use_api_base=True, + auth_kind="web", + ) + if response.status_code == 200: + res_json = response.json() + if res_json.get("code") == 0: + data = res_json.get("data", {}) + items = data.get("items", []) + print(f"Found {len(items)} result(s) for mode '{data.get('mode', mode)}':") + for i, item in enumerate(items, 1): + print(f" [{i}] doc_id: {item.get('doc_id', '')} score: {item.get('score', 0):.4f}") + else: + print(f"Search failed, code: {res_json.get('code')}, message: {res_json.get('message')}") + else: + print(f"Search failed, code: {response.status_code}, body: {response.text}") + def remove_tags(self, command_dict): if self.server_type != "user": print("This command is only allowed in USER mode") @@ -1868,7 +1937,7 @@ class RAGFlowClient: payload["page"] = command_dict["page"] if "size" in command_dict: payload["size"] = command_dict["size"] - if "keywords" in command_dict and command_dict["keywords"]: + if command_dict.get("keywords"): payload["keywords"] = command_dict["keywords"] if "available_int" in command_dict: payload["available_int"] = command_dict["available_int"] @@ -2036,8 +2105,7 @@ class RAGFlowClient: max_width = get_string_width(str(col)) for item in data: value_len = get_string_width(str(item.get(col, ""))) - if value_len > max_width: - max_width = value_len + max_width = max(max_width, value_len) col_widths[col] = max(2, max_width) # Generate delimiter @@ -2230,6 +2298,10 @@ def run_command(client: RAGFlowClient, command_dict: dict): return client.update_chunk(command_dict) case "set_metadata": return client.set_metadata(command_dict) + case "generate_nav_for_dataset": + return client.generate_nav_for_dataset(command_dict) + case "navigation_search": + return client.navigation_search(command_dict) case "remove_tags": return client.remove_tags(command_dict) case "remove_chunks": @@ -2291,6 +2363,8 @@ User Commands (use -t user): LIST DATASETS LIST DOCUMENTS OF DATASET SEARCH ON DATASETS +GENERATE NAVIGATION OF DATASET '' +NAVIGATION SEARCH '' IN DATASET '' MODE '' [TOP_K ] LIST METADATA OF DATASETS [, ]* LIST METADATA SUMMARY OF DATASET DOCUMENTS [, ]* GET CHUNK @@ -2332,9 +2406,8 @@ def run_benchmark(client: RAGFlowClient, command_dict: dict): qps = iterations / total_duration if total_duration > 0 else None print(f"command: {command}, Concurrency: {concurrency}, iterations: {iterations}") print(f"total duration: {total_duration:.4f}s, QPS: {qps}, COMMAND_COUNT: {iterations}, SUCCESS: {success_count}, FAILURE: {iterations - success_count}") - pass else: - results: List[Optional[dict]] = [None] * concurrency + results: list[dict | None] = [None] * concurrency mp_context = mp.get_context("spawn") start_time = time.perf_counter() with ProcessPoolExecutor(max_workers=concurrency, mp_context=mp_context) as executor: @@ -2362,5 +2435,3 @@ def run_benchmark(client: RAGFlowClient, command_dict: dict): qps = total_command_count / total_duration if total_duration > 0 else None print(f"command: {command}, Concurrency: {concurrency} , iterations: {iterations}") print(f"total duration: {total_duration:.4f}s, QPS: {qps}, COMMAND_COUNT: {total_command_count}, SUCCESS: {success_count}, FAILURE: {total_command_count - success_count}") - - pass diff --git a/api/apps/restful_apis/dataset_api.py b/api/apps/restful_apis/dataset_api.py index ccd79fd095..b19d74dd36 100644 --- a/api/apps/restful_apis/dataset_api.py +++ b/api/apps/restful_apis/dataset_api.py @@ -16,10 +16,11 @@ import logging from peewee import OperationalError -from quart import request, make_response -from common.constants import RetCode -from api.apps import login_required, current_user -from api.utils.api_utils import get_error_argument_result, get_error_data_result, get_json_result, get_result, add_tenant_id_to_kwargs +from quart import make_response, request + +from api.apps import current_user, login_required +from api.apps.services import dataset_api_service +from api.utils.api_utils import add_tenant_id_to_kwargs, get_error_argument_result, get_error_data_result, get_json_result, get_result from api.utils.pagination_utils import DEFAULT_PAGE, DEFAULT_PAGE_SIZE, validate_rest_api_ids, validate_rest_api_page, validate_rest_api_page_size from api.utils.validation_utils import ( CreateDatasetReq, @@ -31,7 +32,7 @@ from api.utils.validation_utils import ( validate_and_parse_json_request, validate_and_parse_request_args, ) -from api.apps.services import dataset_api_service +from common.constants import RetCode @manager.route("/datasets/tags/aggregation", methods=["GET"]) # noqa: F821 @@ -980,6 +981,51 @@ async def list_dataset_nav(tenant_id, dataset_id): return get_error_data_result(message="Internal server error") +@manager.route("/datasets//navigation/search", methods=["GET"]) # noqa: F821 +@login_required +@add_tenant_id_to_kwargs +async def search_dataset_nav(tenant_id, dataset_id): + """Unified navigation search across different knowledge layers. + + GET /api/v1/datasets//navigation/search?q=&mode=&top_k=20 + + Modes: + - nav_doc: navigation tree document leaves (default) + - nav_cluster: navigation tree cluster nodes + - navigation_tree: tree-structured BFS beam descent + - chunk: raw document chunks (deduplicated by doc_id) + - all: union of all modes above + + Success: {"code": 0, "data": {"mode": , "total": , + "items": [{"doc_id": str, "score": float}, ...]}} + """ + q = (request.args.get("q") or "").strip() + if not q: + return get_result(data={"mode": request.args.get("mode", "nav_doc"), "total": 0, "items": []}) + mode = (request.args.get("mode") or "nav_doc").strip() + top_k_raw = request.args.get("top_k") + top_k = None + if top_k_raw: + try: + top_k = max(1, int(top_k_raw)) + except (ValueError, TypeError): + return get_error_data_result(message="top_k must be a positive integer") + try: + success, result = await dataset_api_service.search_dataset_layers( + dataset_id, + tenant_id, + q, + mode, + top_k=top_k, + ) + if success: + return get_result(data=result) + return get_result(data=False, message=result, code=RetCode.AUTHENTICATION_ERROR) + except Exception as e: + logging.exception(e) + return get_error_data_result(message="Internal server error") + + @manager.route("/datasets//navigation//children", methods=["GET"]) # noqa: F821 @login_required @add_tenant_id_to_kwargs @@ -1048,6 +1094,38 @@ async def delete_dataset_nav_node(tenant_id, dataset_id, name): return get_error_data_result(message="Internal server error") +@manager.route("/datasets//navigation", methods=["POST"]) # noqa: F821 +@login_required +@add_tenant_id_to_kwargs +async def generate_dataset_nav(tenant_id, dataset_id): + """Create the entire navigation tree from all dataset documents. + + Deletes any existing navigation tree first, then rebuilds it from + scratch. When ``documents`` is provided, only those docs are used; + otherwise all docs in the dataset are auto-discovered. + + POST /api/v1/datasets//navigation + Body (optional): {"documents": [{"doc_id": "...", "summary": "...", + "doc_title": "... (optional)", "source_type": "... (optional)"}, ...]} + Success: {"code": 0, "data": {"deleted": , "upserted": }} + """ + try: + req = await request.json or {} + documents = req.get("documents") + + success, result = await dataset_api_service.generate_nav( + dataset_id, + tenant_id, + documents, + ) + if success: + return get_result(data=result) + return get_result(data=False, message=result, code=RetCode.AUTHENTICATION_ERROR) + except Exception as e: + logging.exception(e) + return get_error_data_result(message="Internal server error") + + @manager.route("/datasets//skills/", methods=["DELETE"]) # noqa: F821 @login_required @add_tenant_id_to_kwargs diff --git a/api/apps/services/dataset_api_service.py b/api/apps/services/dataset_api_service.py index d7a8b85620..b6eff116d4 100644 --- a/api/apps/services/dataset_api_service.py +++ b/api/apps/services/dataset_api_service.py @@ -13,25 +13,24 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import logging import json +import logging import os import re -from api.db.joint_services.tenant_model_service import resolve_model_config, resolve_model_id, get_composite_model_name_by_ids -from common.constants import PAGERANK_FLD, LLMType -from common import settings from api.db.db_models import File +from api.db.joint_services.tenant_model_service import get_composite_model_name_by_ids, resolve_model_config, resolve_model_id +from api.db.services.connector_service import Connector2KbService from api.db.services.document_service import DocumentService, queue_raptor_o_graphrag_tasks from api.db.services.file2document_service import File2DocumentService from api.db.services.file_service import FileService from api.db.services.knowledgebase_service import KnowledgebaseService, validate_dataset_embedding_models -from api.db.services.connector_service import Connector2KbService from api.db.services.task_service import GRAPH_RAPTOR_FAKE_DOC_ID, TaskService from api.db.services.tenant_model_service import TenantModelService from api.db.services.user_service import TenantService, UserService, UserTenantService -from common.constants import FileSource, StatusEnum from api.utils.api_utils import deep_merge, get_parser_config, remap_dictionary_keys, verify_embedding_availability +from common import settings +from common.constants import PAGERANK_FLD, FileSource, LLMType, StatusEnum from common.misc_utils import thread_pool_exec, thread_pool_exec_long_time from rag.advanced_rag.knowlege_compile.wiki import WIKI_PAGE_COMPILE_KWD @@ -382,8 +381,7 @@ async def update_dataset(tenant_id: str, dataset_id: str, req: dict): from rag.nlp import search settings.docStoreConn.update({"exists": PAGERANK_FLD}, {"remove": PAGERANK_FLD}, search.index_name(kb.tenant_id), kb.id) - if "parse_type" in req: - del req["parse_type"] + req.pop("parse_type", None) if not KnowledgebaseService.update_by_id(kb.id, req): return False, "Update dataset error.(Database error)" @@ -547,8 +545,8 @@ def delete_knowledge_graph(dataset_id: str, tenant_id: str): if not KnowledgebaseService.accessible(dataset_id, tenant_id): return False, "no authorization" _, kb = KnowledgebaseService.get_by_id(dataset_id) - from rag.nlp import search from rag.graphrag.phase_markers import clear_phase_markers + from rag.nlp import search settings.docStoreConn.delete({"knowledge_graph_kwd": ["graph", "subgraph", "entity", "relation", "community_report"]}, search.index_name(kb.tenant_id), dataset_id) # Wiping the graph invalidates any phase-completion markers used to @@ -934,8 +932,8 @@ def delete_index(dataset_id: str, tenant_id: str, index_type: str, wipe: bool = TaskService.delete_by_id(task_id) if wipe and index_type == "graph": - from rag.nlp import search from rag.graphrag.phase_markers import clear_phase_markers + from rag.nlp import search settings.docStoreConn.delete({"knowledge_graph_kwd": ["graph", "subgraph", "entity", "relation", "community_report"]}, search.index_name(kb.tenant_id), dataset_id) # Wiping the graph invalidates any phase-completion markers used to @@ -1173,12 +1171,11 @@ def check_embedding(dataset_id: str, tenant_id: str, req: dict): import random import numpy as np - from common.constants import RetCode - from common.doc_store.doc_store_base import OrderByExpr - from rag.nlp import search from api.db.services.llm_service import LLMBundle - from common.constants import LLMType + from common.constants import LLMType, RetCode + from common.doc_store.doc_store_base import OrderByExpr + from rag.nlp import search def _guess_vec_field(src: dict): for k in src or {}: @@ -1843,10 +1840,10 @@ async def get_dataset_structure(dataset_id: str, tenant_id: str, kind: str, keyw return True, empty index_nm, _ = pack - from common.doc_store.doc_store_base import OrderByExpr + from api.apps.services import structure_graph_common as sgc from api.db.services.compilation_template_service import CompilationTemplateService from api.db.services.tenant_llm_service import TenantLLMService - from api.apps.services import structure_graph_common as sgc + from common.doc_store.doc_store_base import OrderByExpr keywords = (keywords or "").strip() _, active_doc_ids = await _current_dataset_docs(dataset_id) @@ -3125,7 +3122,10 @@ def _nav_item(row: dict) -> dict: return { "name": row.get("name") or "", "description": payload.get("description") or "", - # doc_id count under this node: the cluster's tally, or 1 for a leaf. + "keywords": list(payload.get("keywords") or []), + "entities": list(payload.get("entities") or []), + "graph_content": payload.get("graph_content") or "", + # doc_id count under this node: the cluster`s tally, or 1 for a leaf. "doc_count": int(row.get("doc_count_int") or 0) if is_cluster else 1, "type": "cluster" if is_cluster else "doc", "doc_id": None if is_cluster else (row.get("doc_id") or row.get("name")), @@ -3348,6 +3348,450 @@ async def delete_nav_node(dataset_id: str, tenant_id: str, name: str): logging.exception("delete_nav_node: lock release failed for kb=%s", dataset_id) +async def generate_nav( + dataset_id: str, + tenant_id: str, + documents: list[dict] | None = None, +): + """Create the entire navigation tree. + + Deletes any existing navigation tree first, then rebuilds it from + scratch. When ``documents`` is provided, only those doc→summary pairs + are used; otherwise all documents in the dataset are auto-discovered + and inserted into the tree. + + ``documents`` is a list of ``{"doc_id": str, "summary": str, + "doc_title": str (optional), "source_type": str (optional)}``. + + Returns ``(True, {"deleted": , "upserted": })`` on success. + """ + if not KnowledgebaseService.accessible(dataset_id, tenant_id): + return False, "no authorization" + + _, kb = KnowledgebaseService.get_by_id(dataset_id) + if kb is None: + return False, "Dataset not found." + + # Resolve models. + from api.db.joint_services.tenant_model_service import ( + get_tenant_default_model_by_type, + ) + from api.db.services.llm_service import LLMBundle + from common.constants import LLMType + from rag.advanced_rag.knowlege_compile.dataset_nav import upsert_dataset_nav_doc + + if kb.embd_id: + embd_model_config = resolve_model_config(kb.tenant_id, LLMType.EMBEDDING, kb.embd_id) + else: + embd_model_config = get_tenant_default_model_by_type(kb.tenant_id, LLMType.EMBEDDING) + embd_mdl = LLMBundle(kb.tenant_id, embd_model_config) + + chat_model_config = get_tenant_default_model_by_type(kb.tenant_id, LLMType.CHAT) + chat_mdl = LLMBundle(kb.tenant_id, chat_model_config) + + # Step 0: auto-discover documents when not explicitly provided. + if not documents: + try: + # Query all documents in the dataset directly without + # File / File2Document JOINs so that every document + # participates in the navigation tree (including docs + # created via API that have no file record). + from api.db.db_utils import DB + from api.db.services.doc_metadata_service import DocMetadataService + + with DB.connection_context(): + doc_rows = list( + DocumentService.model.select( + DocumentService.model.id, + DocumentService.model.name, + ) + .where( + DocumentService.model.kb_id == dataset_id, + ) + .dicts() + ) + doc_ids = [str(row["id"]) for row in doc_rows] + metadata_map = DocMetadataService.get_metadata_for_documents(doc_ids, dataset_id) if doc_ids else {} + all_docs = [] + for row in doc_rows: + did = str(row["id"]) + all_docs.append( + { + "id": did, + "name": (row.get("name") or ""), + "meta_fields": metadata_map.get(did, {}), + } + ) + + # Prefer RAPTOR-generated summaries stored in the knowledge + # graph (compile_kwd="tree", knowledge_graph_kwd="graph") + # so that rebuilt nav descriptions match the original + # tree-compilation output. Fall back to meta_fields.title + # or filename when the graph is absent. + raptor_summaries: dict[str, str] = {} + pack = _compiled_index_or_none(kb.tenant_id, dataset_id) + if pack is not None: + try: + index_nm, _ = pack + from common.doc_store.doc_store_base import OrderByExpr + + graph_res = settings.docStoreConn.search( + select_fields=["doc_id", "content_with_weight"], + highlight_fields=[], + condition={"compile_kwd": ["tree"], "knowledge_graph_kwd": ["graph"]}, + match_expressions=[], + order_by=OrderByExpr(), + offset=0, + limit=10000, + index_names=index_nm, + knowledgebase_ids=[dataset_id], + ) + graph_map = settings.docStoreConn.get_fields(graph_res, ["doc_id", "content_with_weight"]) + for row in (graph_map or {}).values(): + gid = str(row.get("doc_id") or "") + if not gid: + continue + try: + graph = json.loads(row.get("content_with_weight") or "{}") + except Exception: + continue + entities = graph.get("entities") or [] + relations = graph.get("relations") or [] + child_names = {r.get("to") for r in relations if isinstance(r, dict)} + + # Build both: + # - root_summary: first line of root desc (short, + # for the display title / description field) + # - graph_text: structured text from ALL entities + # and relations (for embedding, keyword extraction, + # entity extraction, and stored as graph_content). + # entity name -> description + name_desc: dict[str, str] = {} + for ent in entities: + if not isinstance(ent, dict): + continue + nm = (ent.get("name") or "").strip() + if nm: + name_desc[nm] = (ent.get("description") or "").strip() + + # root entity = entity whose name never appears as a + # relation target (not a child of anyone). + root_name = "" + root_summary = "" + for ent in entities: + if isinstance(ent, dict) and ent.get("name") not in child_names: + root_name = (ent.get("name") or "").strip() + root_summary = name_desc.get(root_name, "") + root_summary = root_summary.splitlines()[0].strip() if root_summary else root_name + break + + # Build the full graph text + graph_parts: list[str] = [] + if root_name and name_desc.get(root_name): + graph_parts.append(root_name) + graph_parts.append(name_desc[root_name]) + + child_names_set = {n for n in name_desc if n in child_names} + if child_names_set: + graph_parts.append("") + for cname in sorted(child_names_set): + cdesc = name_desc.get(cname, "") + line = f"- {cname}" + if cdesc: + line += f": {cdesc.splitlines()[0].strip()}" + graph_parts.append(line) + + graph_text = "\n".join(graph_parts) if graph_parts else "" + + if root_summary: + raptor_summaries[gid] = { + "title": root_summary, + "graph_text": graph_text or root_summary, + } + except Exception: + logging.exception("generate_nav: failed to read RAPTOR graph summaries for kb=%s", dataset_id) + + documents = [] + for d in all_docs: + doc_id = str(d.get("id", "")) + if not doc_id: + continue + if doc_id in raptor_summaries: + summary = raptor_summaries[doc_id] + else: + meta = d.get("meta_fields") or {} + summary = (meta.get("title") or "").strip() or d.get("name", "") + documents.append( + { + "doc_id": doc_id, + "summary": summary, + } + ) + + except Exception: + logging.exception("generate_nav: failed to auto-discover docs for kb=%s", dataset_id) + return False, "Failed to auto-discover documents." + + if not documents: + return False, "No documents found in dataset." + + # Step 1: delete the entire existing navigation tree so we start clean. + deleted = 0 + pack = _compiled_index_or_none(kb.tenant_id, dataset_id) + if pack is not None: + index_nm, _ = pack + try: + deleted = await thread_pool_exec( + settings.docStoreConn.delete, + {"compile_kwd": [_NAV_COMPILE_KWD]}, + index_nm, + dataset_id, + ) + deleted = int(deleted or 0) + except Exception: + logging.exception("generate_nav: failed to clear existing nav for kb=%s", dataset_id) + return False, "Failed to clear existing navigation tree." + + # Step 2: rebuild the tree from the provided doc→summary pairs. + upserted = 0 + for doc in documents or []: + doc_id = (doc.get("doc_id") or "").strip() + summary = doc.get("summary") + # summary can be a plain string or a RAPTOR tree dict. + if isinstance(summary, str): + summary = summary.strip() + if not doc_id or not summary: + continue + try: + await upsert_dataset_nav_doc( + tenant_id=tenant_id, + kb_id=dataset_id, + doc_id=doc_id, + summary_or_tree=summary, + embd_mdl=embd_mdl, + chat_mdl=chat_mdl, + ) + upserted += 1 + except Exception: + logging.exception("generate_nav: failed for doc=%s kb=%s", doc_id, dataset_id) + return True, {"deleted": deleted, "upserted": upserted, "failed_doc_id": doc_id} + + return True, {"deleted": deleted, "upserted": upserted} + + +# --------------------------------------------------------------------------- +# Unified Explore Search +# --------------------------------------------------------------------------- + +_LAYERS_HANDLERS: dict[str, str] = { + "nav_doc": "_search_layers_nav_docs", + "nav_cluster": "_search_layers_nav_clusters", + "navigation_tree": "_search_layers_navigation_tree", + "chunk": "_search_layers_chunks", + "all": "_search_layers_all", +} + + +async def search_dataset_layers( + dataset_id: str, + tenant_id: str, + query: str, + mode: str, + *, + top_k: int | None = None, +) -> tuple[bool, dict]: + """Unified search across different knowledge layers of a dataset. + + Args: + mode: One of ``"chunk"``, ``"nav_doc"``, ``"nav_cluster"``, ``"navigation_tree"``, ``"all"``. + - chunk: raw document chunks (via the main retrieval pipeline) + - nav_doc: navigation tree document leaves + - nav_cluster: navigation tree cluster nodes + - navigation_tree: tree-structured BFS beam descent + - all: union of all modes, deduplicated by doc_id with best score + + Items are shaped as ``{"doc_id": str, "score": float}``. + """ + from rag.advanced_rag.knowlege_compile.dataset_nav import search_dataset_nav + + if not KnowledgebaseService.accessible(dataset_id, tenant_id): + return False, "no authorization" + if mode not in _LAYERS_HANDLERS: + return False, f"unknown mode: {mode}, expected one of {list(_LAYERS_HANDLERS.keys())}" + _, kb = KnowledgebaseService.get_by_id(dataset_id) + + try: + from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type + from api.db.services.llm_service import LLMBundle + + if kb.embd_id: + embd_model_config = resolve_model_config(kb.tenant_id, LLMType.EMBEDDING, kb.embd_id) + else: + embd_model_config = get_tenant_default_model_by_type(kb.tenant_id, LLMType.EMBEDDING) + embd_mdl = LLMBundle(kb.tenant_id, embd_model_config) + except Exception as e: + logging.warning( + "search_dataset_layers: failed to create LLMBundle(EMBEDDING) for tenant=%s: %s: %s", + kb.tenant_id, + type(e).__name__, + e, + ) + logging.exception("Full traceback for LLMBundle(EMBEDDING) failure") + embd_mdl = None + + if mode == "nav_doc": + return await _search_layers_nav_docs(tenant_id, dataset_id, query, top_k, embd_mdl, search_dataset_nav) + elif mode == "nav_cluster": + return await _search_layers_nav_clusters(tenant_id, dataset_id, query, top_k, embd_mdl, search_dataset_nav) + elif mode == "navigation_tree": + return await _search_layers_navigation_tree(tenant_id, dataset_id, query, top_k, embd_mdl, search_dataset_nav) + elif mode == "chunk": + return await _search_layers_chunks(tenant_id, dataset_id, query, top_k, embd_mdl, kb) + elif mode == "all": + return await _search_layers_all(tenant_id, dataset_id, query, top_k, embd_mdl, kb, search_dataset_nav) + else: + return False, f"unknown mode: {mode}" + + +async def _search_layers_nav_docs(tenant_id, dataset_id, query, top_k, embd_mdl, search_fn): + items = await _nav_search_result( + tenant_id, + dataset_id, + query, + top_k, + embd_mdl, + search_fn, + type_kwd="nav_doc", + ) + return True, {"mode": "nav_doc", "total": len(items), "items": items} + + +async def _search_layers_nav_clusters(tenant_id, dataset_id, query, top_k, embd_mdl, search_fn): + items = await _nav_search_result( + tenant_id, + dataset_id, + query, + top_k, + embd_mdl, + search_fn, + type_kwd="nav_cluster", + ) + return True, {"mode": "nav_cluster", "total": len(items), "items": items} + + +async def _search_layers_navigation_tree(tenant_id, dataset_id, query, top_k, embd_mdl, search_fn): + from rag.advanced_rag.knowlege_compile.dataset_nav import search_nav_tree_descent + + items = await search_nav_tree_descent( + tenant_id, + dataset_id, + query, + embd_mdl, + top_k=top_k, + ) + return True, {"mode": "navigation_tree", "total": len(items), "items": items} + + +async def _nav_search_result(tenant_id, dataset_id, query, top_k, embd_mdl, search_fn, **kwargs): + results = await search_fn( + tenant_id, + dataset_id, + query, + embd_mdl=embd_mdl, + top_k=top_k, + **kwargs, + ) + items: list[dict] = [] + for r in results: + doc_id = (r.get("doc_id") or "").strip() if isinstance(r.get("doc_id"), str) else (r.get("doc_ids") or [None])[0] + items.append( + { + "doc_id": str(doc_id) if doc_id else "", + "score": round(float(r.get("score", 0.0)), 4), + } + ) + return items + + +async def _search_layers_chunks(tenant_id, dataset_id, query, top_k, embd_mdl, kb): + from common import settings + + tenant_ids = [tenant_id] + + kwargs = {} + if top_k is not None: + kwargs["top"] = top_k + + fetch_k = max(top_k, 10) * 3 if top_k is not None else 1024 + try: + ranks = await settings.retriever.retrieval( + query, + embd_mdl, + tenant_ids, + [dataset_id], + 1, + fetch_k, + 0.0, + 0.3, + **kwargs, + ) + except Exception: + return False, "chunk retrieval failed" + + doc_scores: dict[str, float] = {} + for c in ranks.get("chunks", []): + doc_id = (c.get("doc_id") or "").strip() + score = float(c.get("similarity") or c.get("score") or 0.0) + if doc_id and score > doc_scores.get(doc_id, -1.0): + doc_scores[doc_id] = score + + items = sorted( + ({"doc_id": d, "score": round(s, 4)} for d, s in doc_scores.items()), + key=lambda x: x["score"], + reverse=True, + ) + if top_k is not None and top_k > 0: + items = items[:top_k] + + return True, {"mode": "chunk", "total": len(items), "items": items} + + +async def _search_layers_all(tenant_id, dataset_id, query, top_k, embd_mdl, kb, search_fn): + """Run all modes and return the union of doc_ids, with best score per doc.""" + import asyncio as _asyncio + + result_lists = await _asyncio.gather( + _search_layers_nav_docs(tenant_id, dataset_id, query, top_k, embd_mdl, search_fn), + _search_layers_nav_clusters(tenant_id, dataset_id, query, top_k, embd_mdl, search_fn), + _search_layers_navigation_tree(tenant_id, dataset_id, query, top_k, embd_mdl, search_fn), + _search_layers_chunks(tenant_id, dataset_id, query, top_k, embd_mdl, kb), + return_exceptions=True, + ) + + doc_scores: dict[str, float] = {} + for result in result_lists: + if isinstance(result, Exception): + continue + ok, data = result + if not ok: + continue + for item in data.get("items", []): + doc_id = item.get("doc_id", "") + score = float(item.get("score", 0.0)) + if doc_id and score > doc_scores.get(doc_id, -1.0): + doc_scores[doc_id] = score + + items = sorted( + ({"doc_id": d, "score": round(s, 4)} for d, s in doc_scores.items()), + key=lambda x: x["score"], + reverse=True, + ) + if top_k is not None and top_k > 0: + items = items[:top_k] + + return True, {"mode": "all", "total": len(items), "items": items} + + async def update_wiki_page( dataset_id: str, tenant_id: str, @@ -3390,11 +3834,11 @@ async def update_wiki_page( return True, None index_nm, _ = pack - from rag.advanced_rag.knowlege_compile.wiki import ( - _wiki_transform_links, - _wiki_extract_summary, - ) from api.db.services.file_commit_service import FileCommitService + from rag.advanced_rag.knowlege_compile.wiki import ( + _wiki_extract_summary, + _wiki_transform_links, + ) full_slug = f"{page_type}/{slug}" if "/" not in slug else slug diff --git a/rag/advanced_rag/harness/tools/navigation.py b/rag/advanced_rag/harness/tools/navigation.py index 1c03be3cdc..f56915930e 100644 --- a/rag/advanced_rag/harness/tools/navigation.py +++ b/rag/advanced_rag/harness/tools/navigation.py @@ -397,7 +397,13 @@ async def _ask_nav_select(tools, query: str, items: list[dict], noun: str, max_i name = str(it.get("name") or "").strip() or f"item-{i}" desc = str(it.get("description") or "").strip().replace("\n", " ") extra = f" [{it['doc_count']} docs]" if it.get("doc_count") else "" - lines.append(f"[{i}] {name}{extra}: {desc[:300]}") + kwds = it.get("keywords") or [] + tags = ", ".join(str(k) for k in kwds[:6]).strip() + head = f" [tags: {tags}]" if tags else "" + entities = it.get("entities") or [] + ents = ", ".join(str(e) for e in entities[:6]).strip() + head += f" [entities: {ents}]" if ents else "" + lines.append(f"[{i}] {name}{extra}{head}: {desc[:300]}") system = _NAV_SELECT_SYSTEM.format(noun=noun) user = f"Question:\n{query}\n\n{noun.capitalize()} (numbered):\n" + "\n".join(lines) + "\n\nOutput JSON:" diff --git a/rag/advanced_rag/knowlege_compile/dataset_nav.py b/rag/advanced_rag/knowlege_compile/dataset_nav.py index 1d910a802a..e1d19f78b7 100644 --- a/rag/advanced_rag/knowlege_compile/dataset_nav.py +++ b/rag/advanced_rag/knowlege_compile/dataset_nav.py @@ -68,6 +68,54 @@ _LOCK_BLOCKING_TIMEOUT_S = 5 # Hard limit on how many sibling clusters we evaluate per KNN call _KNN_TOP_K = 5 +# Fields needed to shape a nav hit for hybrid scoring / rendering. +_NAV_SEARCH_FIELDS = [ + "id", + "content_with_weight", + "name", + "doc_id", + "type_kwd", + "doc_ids_kwd", + "doc_count_int", +] + +# Weight of the dense leg in hybrid search (1 - this = BM25 leg weight). +_NAV_HYBRID_DENSE_W = 0.5 + +# Stop-words skipped when deriving routing tag-words from a summary. +_NAV_STOP_WORDS = { + "the", + "a", + "an", + "and", + "or", + "of", + "to", + "in", + "on", + "for", + "with", + "at", + "is", + "are", + "was", + "were", + "be", + "been", + "being", + "this", + "that", + "these", + "those", + "it", + "its", + "as", + "by", + "from", + "about", + "into", +} + # --------------------------------------------------------------------------- # Helpers @@ -167,6 +215,60 @@ async def _store_search( return list(rows.values()) +async def _store_text_search( + tenant_id: str, + kb_id: str, + query: str, + fields: list[str], + limit: int = 100, + *, + compile_kwd: str = _COMPILE_KWD, + type_kwd: str = "", + extra_filter: dict | None = None, +) -> list[dict]: + """Full-text (BM25) leg over the nav rows' tokenized fields. + + Recalls nav nodes whose ``content_ltks`` / ``content_sm_ltks`` match the + query — the lexical half of hybrid search, complementing the KNN leg for + exact/proper-noun recall (e.g. "关羽", "ImageNet 2012"). + + Args: + compile_kwd: Which compile partition to search within + (default ``_COMPILE_KWD`` = ``"dataset_nav"``). + type_kwd: Optional type filter (``"nav_doc"``, ``"nav_cluster"``, or ``""``). + extra_filter: Additional filter conditions merged into the query. + """ + from common import settings + from common.doc_store.doc_store_base import MatchTextExpr, OrderByExpr + + # Pre-tokenize the query with the same tokenizer used to index content_ltks. + # ES content_ltks uses the "whitespace" analyzer (no stemming), while our + # Python tokenizer applies stemming (e.g. "Christie" → "christi"). + # Without this, raw query terms won't match pre-stemmed tokens in the index. + tokenized_query = _tokenize(query) + + index = _index_name(tenant_id) + filter_condition: dict = {"compile_kwd": [compile_kwd]} + if type_kwd: + filter_condition["type_kwd"] = type_kwd + if extra_filter: + filter_condition.update(extra_filter) + res = await thread_pool_exec( + settings.docStoreConn.search, + fields, + [], + filter_condition, + [MatchTextExpr(["content_ltks", "content_sm_ltks"], tokenized_query, limit)], + OrderByExpr(), + 0, + limit, + index, + [kb_id], + ) + rows = settings.docStoreConn.get_fields(res, fields) if res else {} + return list(rows.values()) + + async def _store_knn( tenant_id: str, kb_id: str, @@ -320,8 +422,27 @@ def _make_nav_doc_row( depth_int: int, embd_mdl=None, embedding: list[float] | None = None, + *, + graph_content: str = "", ) -> dict: - """Build a nav_doc ES/Infinity row dict for a single document leaf node.""" + """Build a nav_doc ES/Infinity row dict for a single document leaf node. + + Args: + summary: Short human-readable title / description for display. + graph_content: Optional full-graph text used for richer keyword / + entity extraction. When set, keywords and entities are derived + from graph_content instead of summary, and the text is stored in + the ``graph_content`` payload field. + """ + kw_text = graph_content or summary + payload = { + "type": "nav_doc", + "description": summary, + "keywords": _nav_keywords(kw_text), + "entities": _nav_entities(kw_text), + } + if graph_content: + payload["graph_content"] = graph_content row: dict = { "id": _nav_doc_id(doc_id), "kb_id": kb_id, @@ -334,9 +455,8 @@ def _make_nav_doc_row( "depth_int": depth_int, "available_int": 0, } - payload = {"type": "nav_doc", "description": summary} row["content_with_weight"] = json.dumps(payload, ensure_ascii=False) - ltks = _tokenize(summary) + ltks = _tokenize(kw_text) row["content_ltks"] = ltks row["content_sm_ltks"] = _fine_tokenize(ltks) if _vector_len(embedding) > 0: @@ -370,7 +490,12 @@ def _make_nav_cluster_row( "doc_count_int": len(doc_ids), "available_int": 0, } - payload = {"type": "nav_cluster", "description": description} + payload = { + "type": "nav_cluster", + "description": description, + "keywords": _nav_keywords(description), + "entities": _nav_entities(description), + } row["content_with_weight"] = json.dumps(payload, ensure_ascii=False) ltks = _tokenize(description) row["content_ltks"] = ltks @@ -395,6 +520,72 @@ def _fine_tokenize(text: str) -> str: return rag_tokenizer.fine_grained_tokenize(text) +def _nav_keywords(summary: str, max_kwds: int = 6) -> list[str]: + """Derive routing tags from a nav summary (no LLM call, zero cost). + + These are the tokenized non-stop terms the model uses to fast-judge node + relevance before reading the full description. + """ + from rag.nlp import rag_tokenizer + + tokens = (rag_tokenizer.tokenize(summary or "") or "").split() + seen: set[str] = set() + out: list[str] = [] + for t in tokens: + t = t.strip() + if len(t) < 2 or t.isdigit() or t.lower() in _NAV_STOP_WORDS or t.lower() in seen: + continue + seen.add(t.lower()) + out.append(t) + if len(out) >= max_kwds: + break + return out + + +def _nav_entities(summary: str, max_entities: int = 6) -> list[str]: + """Extract likely named entities from a nav summary (no LLM call, zero cost). + + Uses a two-pass heuristic: + 1. Regex for English capitalized multi-word sequences (proper nouns). + 2. Tokenizer-based extraction for CJK and other non-Latin text. + + Returns up to ``max_entities`` deduplicated strings. + """ + text = summary or "" + entities: list[str] = [] + seen: set[str] = set() + + # Pass 1: English-like capitalized sequences (e.g. "New York", "Machine Learning") + for m in re.finditer(r"\b([A-Z][a-z]+(?:\s+[A-Z][a-z]+)+)\b", text): + ent = m.group(1).strip() + key = ent.lower() + if key not in _NAV_STOP_WORDS and key not in seen: + seen.add(key) + entities.append(ent) + if len(entities) >= max_entities: + return entities + + # Pass 2: tokenizer-based for CJK / other scripts + from rag.nlp import rag_tokenizer + + tokens = (rag_tokenizer.tokenize(text) or "").split() + for t in tokens: + t = t.strip() + if len(t) < 3 or t.isdigit(): + continue + if t.isascii() and t[0].islower(): + continue # skip lowercase English fragments + key = t.lower() + if key in _NAV_STOP_WORDS or key in seen: + continue + seen.add(key) + entities.append(t) + if len(entities) >= max_entities: + break + + return entities + + def _matches_condition(row: dict, condition: dict) -> bool: """Check the simple equality filters used by dataset navigation.""" for field, expected in condition.items(): @@ -574,17 +765,27 @@ async def upsert_dataset_nav_doc( tenant_id: Tenant owning the KB. kb_id: Knowledge base id. doc_id: Document id. - summary_or_tree: A plain summary string, or a RAPTOR tree dict from - which the root summary is extracted. + summary_or_tree: + - A plain summary string (fallback, no RAPTOR graph). + - A RAPTOR tree dict (backwards-compat, root summary extracted). + - A dict with keys ``title`` (short display label) and + ``graph_text`` (full graph entities + relations for richer + embedding / keyword extraction). embd_mdl: LLMBundle for embedding (required for clustering). chat_mdl: LLMBundle for chat (required for LLM merge/summary). """ if not doc_id or not kb_id: return - # 1. Extract summary + # 1. Extract display summary and optional full graph content. + graph_content = "" if isinstance(summary_or_tree, dict): - summary = _extract_root_summary_from_tree(summary_or_tree) + if "title" in summary_or_tree and "graph_text" in summary_or_tree: + # New extended format from generate_nav's RAPTOR graph path. + summary = (summary_or_tree.get("title") or "").strip() + graph_content = (summary_or_tree.get("graph_text") or "").strip() + else: + summary = _extract_root_summary_from_tree(summary_or_tree) elif isinstance(summary_or_tree, str): summary = summary_or_tree else: @@ -593,10 +794,10 @@ async def upsert_dataset_nav_doc( logging.info("dataset_nav: skipping doc=%s (kb=%s) — no summary", doc_id, kb_id) return - # 2. Embed doc summary before taking the KB lock. The result is - # independent of the nav tree; all tree reads, deletes, and writes below - # are serialized by the same lock. - doc_embedding = await _embed(embd_mdl, summary) if embd_mdl else [] + # 2. Embed with the richest available text so that KNN placement + # benefits from the full tree structure when present. + embed_text = graph_content or summary + doc_embedding = await _embed(embd_mdl, embed_text) if embd_mdl else [] vec_dim = len(doc_embedding) lock = RedisDistributedLock( @@ -665,6 +866,7 @@ async def upsert_dataset_nav_doc( depth, embd_mdl, doc_embedding, + graph_content=graph_content, ) await _store_upsert(tenant_id, kb_id, nav_doc_row) @@ -712,6 +914,7 @@ async def upsert_dataset_nav_doc( new_depth, embd_mdl, doc_embedding, + graph_content=graph_content, ) await _store_upsert(tenant_id, kb_id, nav_doc_row) else: @@ -739,6 +942,7 @@ async def upsert_dataset_nav_doc( 1, embd_mdl, doc_embedding, + graph_content=graph_content, ) await _store_upsert(tenant_id, kb_id, nav_doc_row) @@ -1015,29 +1219,47 @@ async def search_dataset_nav( kb_id: str, query: str, embd_mdl=None, - top_k: int = 8, + top_k: int | None = None, + *, + type_kwd: str = "", + compile_kwd: str = _COMPILE_KWD, ) -> list[dict]: """Find the nav-tree nodes most relevant to ``query`` for one KB. The nav rows are ``available_int=0`` (invisible to the normal retriever), so this is the sanctioned read seam: a caller uses the returned document ids to - route a scoped chunk retrieval. Returns items shaped as:: + route a scoped chunk retrieval. + + Args: + type_kwd: Optional type filter — ``"nav_doc"`` to restrict to document + leaves, ``"nav_cluster"`` to restrict to clusters, or ``""`` for all. + compile_kwd: Which compile partition to search within + (default ``_COMPILE_KWD`` = ``"dataset_nav"``). + + Returns items shaped as:: {"type": "nav_doc" | "nav_cluster", "doc_id": str | None, # the document, for a leaf "doc_ids": [str], # the documents a node covers "name": str, "description": str, "score": float} - Ranked by vector KNN over the node summaries when ``embd_mdl`` is given; - otherwise a best-effort text-ranked scan. + Ranked by hybrid search: KNN over the node-summary vectors fused with a + BM25 (full-text) leg over the tokenized fields. When ``embd_mdl`` is absent + the lexical leg alone ranks the results. """ query = (query or "").strip() if not query: return [] - condition = {"compile_kwd": [_COMPILE_KWD]} - rows_with_scores: list[tuple[dict, float]] = [] + condition: dict = {"compile_kwd": [compile_kwd]} + if type_kwd: + condition["type_kwd"] = type_kwd + # name -> [row, fused_score]; `name` uniquely identifies a nav node, so it + # is the dedup key when the dense and lexical legs return the same node. + fused: dict[str, list] = {} + dense_w = _NAV_HYBRID_DENSE_W if embd_mdl is not None else 0.0 + # ── Dense leg: KNN over the node-summary vectors ── if embd_mdl is not None: try: vec = await _embed(embd_mdl, query) @@ -1046,28 +1268,39 @@ async def search_dataset_nav( vec = [] if _vector_len(vec) > 0: try: - rows = await _store_knn(tenant_id, kb_id, vec, len(vec), condition, top_k=top_k) + rows = await _store_knn(tenant_id, kb_id, vec, len(vec), condition, top_k=top_k or 10000) vf = _vec_field(len(vec)) - rows_with_scores = [(r, _cosine_sim(vec, r.get(vf))) for r in rows] + for r in rows: + rk = r.get("name") or r.get("doc_id") or "" + if not rk: + continue + fused.setdefault(rk, [r, 0.0])[1] += dense_w * _cosine_sim(vec, r.get(vf)) except Exception: logging.exception("search_dataset_nav: knn failed for kb=%s", kb_id) - rows_with_scores = [] - if not rows_with_scores: - fields = ["content_with_weight", "name", "doc_id", "type_kwd", "doc_ids_kwd", "doc_count_int"] - try: - rows = await _store_search(tenant_id, kb_id, condition, fields, limit=max(top_k * 20, 100)) - except Exception: - logging.exception("search_dataset_nav: scan failed for kb=%s", kb_id) - rows = [] - rows_with_scores = [(r, _nav_text_score(query, r)) for r in rows] - rows_with_scores.sort(key=lambda item: item[1], reverse=True) + # ── Lexical leg: engine BM25 over the tokenized fields ── + text_w = 1.0 - dense_w + try: + text_rows = await _store_text_search(tenant_id, kb_id, query, _NAV_SEARCH_FIELDS, limit=max((top_k or 0) * 3, 20) if top_k else 10000, compile_kwd=compile_kwd, type_kwd=type_kwd) + except Exception: + logging.exception("search_dataset_nav: text search failed for kb=%s", kb_id) + text_rows = [] + for r in text_rows: + rk = r.get("name") or r.get("doc_id") or "" + if not rk: + continue + ts = _nav_text_score(query, r) + if ts <= 0: + continue + fused.setdefault(rk, [r, 0.0])[1] += text_w * ts - # Discard zero-score rows (text match produced no relevant hits) - rows_with_scores = [(r, s) for r, s in rows_with_scores if s > 0] + rows_with_scores = [(r, s) for r, s in fused.values() if s > 0] + rows_with_scores.sort(key=lambda item: item[1], reverse=True) + if top_k is not None: + rows_with_scores = rows_with_scores[:top_k] out: list[dict] = [] - for r, score in rows_with_scores[:top_k]: + for r, score in rows_with_scores: try: payload = json.loads(r.get("content_with_weight") or "{}") except Exception: @@ -1088,6 +1321,11 @@ async def search_dataset_nav( "doc_ids": doc_ids, "name": name, "description": payload.get("description") or "", + "keywords": _as_str_list(payload.get("keywords")), + "entities": _as_str_list(payload.get("entities")), + "graph_content": payload.get("graph_content") or "", + "doc_title": payload.get("doc_title") or "", + "source_type": payload.get("source_type") or "", "doc_count": int(r.get("doc_count_int") or len(doc_ids) or 0), "score": float(score or 0.0), } @@ -1095,6 +1333,233 @@ async def search_dataset_nav( return out +async def search_nav_tree_descent( + tenant_id: str, + kb_id: str, + query: str, + embd_mdl, + top_k: int | None = None, +) -> list[dict]: + """Tree-structured hybrid search: descend from root into the most relevant branches. + + Unlike ``search_dataset_nav`` (flat hybrid search over all nav nodes), + this respects the parent-child hierarchy: it starts at root clusters, + then at each level descends into the semantically closest children, + collecting document identifiers from ``nav_doc`` leaves along the way. + + Each level uses hybrid search (KNN + BM25 text) fused with the same + weights as ``search_dataset_nav``, so both vector similarity *and* + lexical matching drive the descent. + + The search uses BFS with beam pruning — at each depth, only the + *beam_width* most similar clusters are expanded further. + + Returns items shaped as ``{"doc_id": str, "score": float}``. + """ + query = (query or "").strip() + if not query: + return [] + if embd_mdl is None: + logging.warning( + "search_nav_tree_descent: embd_mdl is None — falling back to text-only flat search for kb=%s query=%.80s", + kb_id, + query, + ) + raw = await search_dataset_nav(tenant_id, kb_id, query, embd_mdl=None, top_k=top_k, type_kwd="nav_doc") + return [{"doc_id": r.get("doc_id", ""), "score": r.get("score", 0.0)} for r in raw if r.get("doc_id")] + + vec = await _embed(embd_mdl, query) + vec_dim = _vector_len(vec) + if vec_dim == 0: + return [] + + beam_width = 5 + dense_w = _NAV_HYBRID_DENSE_W # same weight as search_dataset_nav + vf = _vec_field(vec_dim) + + fields = [ + "content_with_weight", + "name", + "doc_id", + "compile_kwd", + "type_kwd", + "parent_kwd", + "depth_int", + "doc_count_int", + "doc_ids_kwd", + vf, + ] + + collected: list[dict] = [] + seen_docs: set[str] = set() + seen_nodes: set[str] = set() + + # ── Root clusters (KNN only — no text) ── + # Root clusters at depth=0 are LLM-generated broad-topic summaries. + # Specific query terms (e.g. "British") almost never appear in their + # name / description, so we route purely by vector similarity at this + # level. Text matching takes over from depth ≥ 1 where cluster + # descriptions contain concrete terms. + root_cond = { + "kb_id": [kb_id], + "compile_kwd": [_COMPILE_KWD], + "type_kwd": ["nav_cluster"], + "depth_int": [0], + } + roots_knn = await _store_knn(tenant_id, kb_id, vec, vec_dim, root_cond, top_k=beam_width * 3) + + # If no root cluster (depth=0) exists — the dataset may have been + # compiled without one — scan all nav_clusters to find the lowest + # available depth and start beam search there. + if not roots_knn: + all_cond = { + "kb_id": [kb_id], + "compile_kwd": [_COMPILE_KWD], + "type_kwd": ["nav_cluster"], + } + all_clusters = await _store_search(tenant_id, kb_id, all_cond, fields, limit=10000) + if not all_clusters: + return [] + + min_depth = min( + (r.get("depth_int", 0) for r in all_clusters if r.get("depth_int") is not None), + default=0, + ) + logging.warning( + "search_nav_tree_descent: no root cluster at depth=0, starting beam search from depth=%d", + min_depth, + ) + + starters = [r for r in all_clusters if r.get("depth_int") == min_depth] + starters.sort(key=lambda r: _cosine_sim(vec, r.get(vf)), reverse=True) + current_level = starters[:beam_width] + for r in current_level: + r["_score"] = _cosine_sim(vec, r.get(vf)) + else: + # Route down to beam_width semantically closest root clusters. + for r in roots_knn: + r["_score"] = _cosine_sim(vec, r.get(vf)) + roots_knn.sort(key=lambda r: r["_score"], reverse=True) + current_level = roots_knn[:beam_width] + + if not current_level: + return [] + + while current_level and (top_k is None or len(collected) < top_k): + next_level: list[dict] = [] + + for node in current_level: + node_name = node.get("name", "") + if node_name in seen_nodes: + continue + seen_nodes.add(node_name) + parent_score = node.get("_score", 0.0) + + child_cond: dict = { + "kb_id": [kb_id], + "compile_kwd": [_COMPILE_KWD], + "parent_kwd": [node_name], + } + children_knn = await _store_knn(tenant_id, kb_id, vec, vec_dim, child_cond, top_k=beam_width * 3) + children_text = await _store_text_search( + tenant_id, + kb_id, + query, + fields, + limit=beam_width * 3, + extra_filter={"parent_kwd": [node_name], "kb_id": [kb_id]}, + ) + candidates = _hybrid_fuse(vec, vf, query, children_knn, children_text, dense_w, beam_width) + + for c in candidates: + if top_k is not None and len(collected) >= top_k: + break + if c.get("type_kwd") == "nav_doc": + doc_id = (c.get("doc_id") or "").strip() + if doc_id and doc_id not in seen_docs: + seen_docs.add(doc_id) + collected.append({"doc_id": doc_id, "score": round(c["_score"] or parent_score, 4)}) + else: + next_level.append(c) + + if top_k is not None and len(collected) >= top_k: + break + + if top_k is not None and len(collected) >= top_k: + break + + next_level.sort(key=lambda c: c.get("_score", 0.0), reverse=True) + current_level = next_level[:beam_width] + + # Fallback: collect doc_ids from terminal cluster nodes. + if not collected and current_level: + for node in current_level: + for did in node.get("doc_ids_kwd") or []: + did_str = str(did).strip() + if did_str and did_str not in seen_docs: + seen_docs.add(did_str) + collected.append({"doc_id": did_str, "score": round(node.get("_score", 0.0), 4)}) + if top_k is not None and len(collected) >= top_k: + break + if top_k is not None and len(collected) >= top_k: + break + + return collected + + +def _hybrid_fuse( + vec: list[float], + vf: str, + query: str, + knn_rows: list[dict], + text_rows: list[dict], + dense_w: float, + top_k: int, +) -> list[dict]: + """Fuse KNN and text results into a scored, deduplicated, sorted list. + + ``_store_knn`` already enforces filter conditions on the KNN leg; + text rows come from ``_store_text_search`` which applies its own + filters. Here we merge both legs by ``name``/``doc_id``. + """ + text_w = 1.0 - dense_w + fused: dict[str, tuple[dict, float]] = {} # key → (row, score) + + # KNN leg: raw cosine * dense_w, matching search_dataset_nav's scale so + # both search paths score the dense leg identically. + for r in knn_rows: + rk = r.get("name") or r.get("doc_id") or "" + if not rk: + continue + key = f"knn:{rk}" + fused[key] = (r, _cosine_sim(vec, r.get(vf, [])) * dense_w) + + # Text leg. + for r in text_rows: + rk = r.get("name") or r.get("doc_id") or "" + if not rk: + continue + ts = _nav_text_score(query, r) + if ts <= 0: + continue + key = f"text:{rk}" + if key in fused: + fused[key] = (fused[key][0], fused[key][1] + text_w * ts) + else: + key_knn = f"knn:{rk}" + if key_knn in fused: + fused[key_knn] = (fused[key_knn][0], fused[key_knn][1] + text_w * ts) + else: + fused[key] = (r, text_w * ts) + + rows_with_scores = [(r, s) for r, s in fused.values() if s > 0] + rows_with_scores.sort(key=lambda item: item[1], reverse=True) + result = rows_with_scores[:top_k] + for r, s in result: + r["_score"] = s + return [r for r, _ in result] + + def _as_str_list(value) -> list[str]: if isinstance(value, list): return [str(v) for v in value if v] @@ -1108,11 +1573,17 @@ def _nav_text_score(query: str, row: dict) -> float: payload = json.loads(row.get("content_with_weight") or "{}") except Exception: payload = {} + keywords = payload.get("keywords") or [] + entities = payload.get("entities") or [] + graph_content = payload.get("graph_content") or "" haystack = " ".join( str(x or "") for x in ( row.get("name"), payload.get("description"), + graph_content, + *keywords, + *entities, ) ).lower() q_terms = set(re.findall(r"[\w]+", query.lower())) diff --git a/rag/utils/es_conn.py b/rag/utils/es_conn.py index b06a2c3e5c..1b15722c72 100644 --- a/rag/utils/es_conn.py +++ b/rag/utils/es_conn.py @@ -14,18 +14,19 @@ # limitations under the License. # -import re +import copy import json +import re import time -import copy -from elasticsearch_dsl import UpdateByQuery, Q, Search from elastic_transport import ConnectionTimeout +from elasticsearch_dsl import Q, Search, UpdateByQuery + +from common.constants import PAGERANK_FLD, TAG_FLD from common.decorator import singleton -from common.doc_store.doc_store_base import MatchTextExpr, OrderByExpr, MatchExpr, MatchDenseExpr, FusionExpr +from common.doc_store.doc_store_base import FusionExpr, MatchDenseExpr, MatchExpr, MatchTextExpr, OrderByExpr from common.doc_store.es_conn_base import ESConnectionBase from common.float_utils import get_float -from common.constants import PAGERANK_FLD, TAG_FLD ATTEMPT_TIME = 2 MAX_RESULT_WINDOW = 10000 @@ -227,7 +228,7 @@ class ESConnection(ESConnectionBase): elif isinstance(v, str) or isinstance(v, int): bool_query.filter.append(Q("term", **{k: v})) else: - raise Exception(f"Condition `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str or list.") + raise Exception(f"Condition `{k!s}={v!s}` value type is {type(v)!s}, expected to be int, str or list.") s = Search() vector_similarity_weight = 0.5 @@ -243,7 +244,7 @@ class ESConnection(ESConnectionBase): vector_similarity_weight = get_float(weights.split(",")[1]) for m in match_expressions: if isinstance(m, MatchTextExpr): - minimum_should_match = m.extra_options.get("minimum_should_match", 0.0) + minimum_should_match = (m.extra_options or {}).get("minimum_should_match", 0.0) if isinstance(minimum_should_match, float): minimum_should_match = str(int(minimum_should_match * 100)) + "%" bool_query.must.append(Q("query_string", fields=m.fields, type="best_fields", query=m.matching_text, minimum_should_match=minimum_should_match, boost=1)) @@ -254,10 +255,11 @@ class ESConnection(ESConnectionBase): similarity = 0.0 if "similarity" in m.extra_options: similarity = m.extra_options["similarity"] + k = min(m.topn, 10000) s = s.knn( m.vector_column_name, - m.topn, - m.topn * 2, + k, + min(k * 2, 10000), query_vector=list(m.embedding_data), filter=bool_query.to_dict(), # filter=_build_knn_filter_query(bool_query, vector_similarity_weight), similarity=similarity, @@ -310,7 +312,7 @@ class ESConnection(ESConnectionBase): vector_fields = [f for f in (select_fields or []) if f.endswith("_vec")] if vector_fields: q["fields"] = vector_fields - self.logger.debug(f"ESConnection.search {str(index_names)} query: " + json.dumps(q)) + self.logger.debug(f"ESConnection.search {index_names!s} query: " + json.dumps(q)) for i in range(ATTEMPT_TIME): try: @@ -321,7 +323,7 @@ class ESConnection(ESConnectionBase): res = self._es_search_once(index_names, q, track_total_hits=True) if str(res.get("timed_out", "")).lower() == "true": raise Exception("Es Timeout.") - self.logger.debug(f"ESConnection.search {str(index_names)} res: " + str(res)) + self.logger.debug(f"ESConnection.search {index_names!s} res: " + str(res)) return res except ConnectionTimeout: self.logger.exception("ES request timeout") @@ -330,9 +332,9 @@ class ESConnection(ESConnectionBase): except Exception as e: # Only log debug for NotFoundError(accepted when metadata index doesn't exist) if "NotFound" in str(e): - self.logger.debug(f"ESConnection.search {str(index_names)} query: " + str(q) + " - " + str(e)) + self.logger.debug(f"ESConnection.search {index_names!s} query: " + str(q) + " - " + str(e)) else: - self.logger.exception(f"ESConnection.search {str(index_names)} query: " + str(q) + str(e)) + self.logger.exception(f"ESConnection.search {index_names!s} query: " + str(q) + str(e)) raise e self.logger.error(f"ESConnection.search timeout for {ATTEMPT_TIME} times!") @@ -443,7 +445,7 @@ class ESConnection(ESConnectionBase): elif isinstance(v, str) or isinstance(v, int): bool_query.filter.append(Q("term", **{k: v})) else: - raise Exception(f"Condition `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str or list.") + raise Exception(f"Condition `{k!s}={v!s}` value type is {type(v)!s}, expected to be int, str or list.") scripts = [] params = {} for k, v in new_value.items(): @@ -473,7 +475,7 @@ class ESConnection(ESConnectionBase): scripts.append(f"ctx._source.{k}=params.pp_{k};") params[f"pp_{k}"] = json.dumps(v, ensure_ascii=False) else: - raise Exception(f"newValue `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str.") + raise Exception(f"newValue `{k!s}={v!s}` value type is {type(v)!s}, expected to be int, str.") ubq = UpdateByQuery(index=index_name).using(self.es).query(bool_query) ubq = ubq.script(source="".join(scripts), params=params) ubq = ubq.params(refresh=True) diff --git a/web/src/interfaces/database/dataset-nav.ts b/web/src/interfaces/database/dataset-nav.ts index de03fa81ce..7916a2b995 100644 --- a/web/src/interfaces/database/dataset-nav.ts +++ b/web/src/interfaces/database/dataset-nav.ts @@ -3,8 +3,11 @@ export interface DatasetNavNode { description: string; doc_count: number; type: string; - doc_id?: string; // only returned by the children endpoint + doc_id?: string; has_children: boolean; + keywords?: string[]; + entities?: string[]; + graph_content?: string; } export interface DatasetNavList { diff --git a/web/src/locales/en.ts b/web/src/locales/en.ts index 2531a8bd52..47f7299052 100644 --- a/web/src/locales/en.ts +++ b/web/src/locales/en.ts @@ -2399,6 +2399,10 @@ Example: Virtual Hosted Style`, loading: 'Loading...', selectNode: 'Select a child node to view details', noDescription: 'No description', + description: 'Description', + keywords: 'Keywords', + entities: 'Entities', + graphContent: 'Full Graph Content', docCount: '{{count}} documents', deleteAllTitle: 'Delete navigation tree', deleteAllDescription: diff --git a/web/src/locales/zh.ts b/web/src/locales/zh.ts index 4f5d86ab23..22e7eb4839 100644 --- a/web/src/locales/zh.ts +++ b/web/src/locales/zh.ts @@ -2045,6 +2045,10 @@ NER：使用 spaCy NER 和基于规则的关键词提取来抽取实体和关系 loading: '加载中...', selectNode: '选择子节点以查看详情', noDescription: '暂无描述', + description: '描述', + keywords: '关键词', + entities: '实体', + graphContent: '完整图谱内容', docCount: '{{count}} 个文档', deleteAllTitle: '删除目录树', deleteAllDescription: '确定要删除整个目录树吗？此操作无法撤销。', diff --git a/web/src/pages/dataset/compilation/hooks/use-compilation-nav.ts b/web/src/pages/dataset/compilation/hooks/use-compilation-nav.ts index 70e5c01537..f3739a0f19 100644 --- a/web/src/pages/dataset/compilation/hooks/use-compilation-nav.ts +++ b/web/src/pages/dataset/compilation/hooks/use-compilation-nav.ts @@ -12,6 +12,9 @@ export interface SelectedNavNode { name: string; description: string; doc_count: number; + keywords?: string[]; + entities?: string[]; + graph_content?: string; } export function useCompilationNav() { @@ -95,6 +98,9 @@ export function useCompilationNav() { name: node.name, description: node.description, doc_count: node.doc_count, + keywords: node.keywords, + entities: node.entities, + graph_content: node.graph_content, }); }, [], diff --git a/web/src/pages/dataset/compilation/nav-tree-view.tsx b/web/src/pages/dataset/compilation/nav-tree-view.tsx index 02868f1f79..0b6414ab17 100644 --- a/web/src/pages/dataset/compilation/nav-tree-view.tsx +++ b/web/src/pages/dataset/compilation/nav-tree-view.tsx @@ -52,8 +52,59 @@ export function NavTreeView() { {t('datasetNav.docCount', { count: selectedNode.doc_count })} -
- {selectedNode.description || t('datasetNav.noDescription')} +
+
+

+ {t('datasetNav.description')} +

+

+ {selectedNode.description || t('datasetNav.noDescription')} +

+
+ {selectedNode.keywords && selectedNode.keywords.length > 0 && ( +
+

+ {t('datasetNav.keywords')} +

+
+ {selectedNode.keywords.map((kw, i) => ( + + {kw} + + ))} +
+
+ )} + {selectedNode.entities && selectedNode.entities.length > 0 && ( +
+

+ {t('datasetNav.entities')} +

+
+ {selectedNode.entities.map((entity, i) => ( + + {entity} + + ))} +
+
+ )} + {selectedNode.graph_content && ( +
+

+ {t('datasetNav.graphContent')} +

+

+ {selectedNode.graph_content} +

+
+ )}
) : (