diff --git a/api/apps/services/dataset_api_service.py b/api/apps/services/dataset_api_service.py index 65ff56e53c..e09a18b61d 100644 --- a/api/apps/services/dataset_api_service.py +++ b/api/apps/services/dataset_api_service.py @@ -33,6 +33,7 @@ from api.db.services.user_service import TenantService, UserService, UserTenantS 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.misc_utils import thread_pool_exec +from rag.advanced_rag.knowlege_compile.wiki import WIKI_PAGE_COMPILE_KWD _VALID_INDEX_TYPES = {"graph", "raptor", "mindmap", "artifact", "skill"} @@ -1481,8 +1482,6 @@ async def search_datasets(tenant_id: str, req: dict): # source_chunk_ids, source_doc_ids # --------------------------------------------------------------------------- -_WIKI_COMPILE_KWD = "artifact_page" -_WIKI_TOPIC_COMPILE_KWD = "artifact_page_topic" _SKILL_COMPILE_KWD = "skill" _SKILL_ALL_COMPILE_KWD = "skill_all" @@ -1528,7 +1527,7 @@ async def has_any_wiki(dataset_id: str, tenant_id: str): res = settings.docStoreConn.search( select_fields=["id"], highlight_fields=[], - condition={"compile_kwd": [_WIKI_COMPILE_KWD]}, + condition={"compile_kwd": [WIKI_PAGE_COMPILE_KWD]}, match_expressions=[], order_by=OrderByExpr(), offset=0, @@ -1575,7 +1574,7 @@ async def list_wiki_pages( page_type = page_type.strip() if isinstance(page_type, str) else page_type topic = topic.strip() if isinstance(topic, str) else topic - condition: dict = {"compile_kwd": [_WIKI_COMPILE_KWD]} + condition: dict = {"compile_kwd": [WIKI_PAGE_COMPILE_KWD]} if page_type: condition["page_type_kwd"] = [page_type] if topic: @@ -1658,50 +1657,75 @@ async def list_wiki_topics( page_size = max(1, min(int(page_size or 200), 1000)) offset = (page - 1) * page_size - order_by = OrderByExpr() try: - order_by.asc("title_kwd") - except Exception: - order_by = OrderByExpr() - - select_fields = [ - "id", - "topic_kwd", - "title_kwd", - "slug_kwd", - ] - try: - res = settings.docStoreConn.search( - select_fields=select_fields, + agg_res = settings.docStoreConn.search( + select_fields=["id"], highlight_fields=[], - condition={"compile_kwd": [_WIKI_TOPIC_COMPILE_KWD]}, + condition={"compile_kwd": [WIKI_PAGE_COMPILE_KWD], "page_type_kwd": ["concept", "entity"]}, match_expressions=[], - order_by=order_by, - offset=offset, - limit=page_size, + order_by=OrderByExpr(), + offset=0, + limit=0, index_names=index_nm, knowledgebase_ids=[dataset_id], + agg_fields=["topic_kwd"], ) - field_map = settings.docStoreConn.get_fields(res, select_fields) + buckets = settings.docStoreConn.get_aggregation(agg_res, "topic_kwd") except Exception: - logging.exception("list_wiki_topics: docStore search failed for kb=%s", dataset_id) + logging.exception("list_wiki_topics: docStore aggregation failed for kb=%s", dataset_id) return True, {"total": 0, "items": []} - total = settings.docStoreConn.get_total(res) - items = [] - for row in (field_map or {}).values(): - topic = row.get("topic_kwd") - if not isinstance(topic, str) or not topic: - continue - items.append( - { - "topic": topic, - "title": row.get("title_kwd") or topic, - "slug": row.get("slug_kwd") or topic, - } - ) + counts = {t: int(c) for t, c in (buckets or []) if isinstance(t, str) and t and int(c or 0) > 0} + if not counts: + return True, {"total": 0, "items": []} - return True, {"total": int(total or 0), "items": items} + # Resolve display metadata (title/slug) from the topic landing pages; fall + # back to the raw topic name when a topic has no ``page_type="topic"`` row. + meta: dict[str, dict] = {} + try: + meta_fields = ["topic_kwd", "title_kwd", "slug_kwd"] + _BATCH = 1000 + _offset = 0 + while True: + meta_res = settings.docStoreConn.search( + select_fields=meta_fields, + highlight_fields=[], + condition={"compile_kwd": [WIKI_PAGE_COMPILE_KWD], "page_type_kwd": ["topic"]}, + match_expressions=[], + order_by=OrderByExpr(), + offset=_offset, + limit=_BATCH, + index_names=index_nm, + knowledgebase_ids=[dataset_id], + ) + rows = settings.docStoreConn.get_fields(meta_res, meta_fields) or {} + if not rows: + break + for row in rows.values(): + t = row.get("topic_kwd") + if isinstance(t, str) and t: + meta[t] = {"title": row.get("title_kwd") or t, "slug": row.get("slug_kwd") or t} + _offset += _BATCH + except Exception: + logging.exception("list_wiki_topics: topic metadata lookup failed for kb=%s", dataset_id) + + # Rank topics by page count (descending), then title for a stable order. + ranked = sorted( + ( + { + "topic": t, + "title": (meta.get(t) or {}).get("title") or t, + "slug": (meta.get(t) or {}).get("slug") or t, + "page_count": c, + } + for t, c in counts.items() + ), + key=lambda x: (-x["page_count"], x["title"].lower()), + ) + + total = len(ranked) + items = ranked[offset : offset + page_size] + return True, {"total": total, "items": items} async def get_wiki_page( @@ -1750,7 +1774,7 @@ async def get_wiki_page( select_fields=select_fields, highlight_fields=[], condition={ - "compile_kwd": [_WIKI_COMPILE_KWD], + "compile_kwd": [WIKI_PAGE_COMPILE_KWD], "page_type_kwd": [page_type], "slug_kwd": [full_slug], }, @@ -2001,7 +2025,7 @@ async def update_wiki_page( select_fields=["id", "content_with_weight"], highlight_fields=[], condition={ - "compile_kwd": [_WIKI_COMPILE_KWD], + "compile_kwd": [WIKI_PAGE_COMPILE_KWD], "page_type_kwd": [page_type], "slug_kwd": [full_slug], }, @@ -2101,7 +2125,7 @@ _WIKI_COMPILE_KWDS = ( "artifact_compilation_plan", "artifact_page_draft", "artifact_page", - _WIKI_TOPIC_COMPILE_KWD, + "artifact_page_topic", "artifact_entity", "artifact_relation", ) diff --git a/rag/advanced_rag/agentic_rag_graph.py b/rag/advanced_rag/agentic_rag_graph.py index ba8a8c0bc7..036b432103 100644 --- a/rag/advanced_rag/agentic_rag_graph.py +++ b/rag/advanced_rag/agentic_rag_graph.py @@ -164,7 +164,7 @@ def build_agentic_graph(tools, token_queue: asyncio.Queue, gen_conf: dict | None q, kw = await tools.formalize(msgs) q = (q or "").strip() kw = (kw or "").strip() - _LOG.info("[Formalizing the question] Understood the question as: \"%s\" — searching with keywords: %s", _snip(q), _snip(kw)) + _LOG.info('[Formalizing the question] Understood the question as: "%s" — searching with keywords: %s', _snip(q), _snip(kw)) return { "question": q, "keywords": kw, @@ -198,7 +198,7 @@ def build_agentic_graph(tools, token_queue: asyncio.Queue, gen_conf: dict | None q = state.get("question", "") kw = state.get("keywords", "") - _LOG.info("[Preliminary search] Taking a first look in the knowledge base for: \"%s\" (keywords: %s)", _snip(q), _snip(kw)) + _LOG.info('[Preliminary search] Taking a first look in the knowledge base for: "%s" (keywords: %s)', _snip(q), _snip(kw)) try: result = await hybrid_search(tools, query=q, keywords=kw) except Exception: @@ -231,7 +231,7 @@ def build_agentic_graph(tools, token_queue: asyncio.Queue, gen_conf: dict | None empty_result = state.get("empty_result", False) _note = " — partial answer, some gaps remain" if partial else (" — not enough evidence to answer" if abstain else "") - _LOG.info("[Composing the answer] Writing the final answer to \"%s\" from %d gathered passage(s)%s.", _snip(question), len(kbinfos["chunks"]), _note) + _LOG.info('[Composing the answer] Writing the final answer to "%s" from %d gathered passage(s)%s.', _snip(question), len(kbinfos["chunks"]), _note) tools.kbinfos = kbinfos @@ -302,7 +302,12 @@ def build_agentic_graph(tools, token_queue: asyncio.Queue, gen_conf: dict | None async def run_agentic_rag(tools, messages: list, max_loops: int = 3, gen_conf: dict | None = None): """Drive the agentic-search graph, yielding answer-token strings.""" - _LOG.info("[Agentic RAG] Starting research to answer your question...") + _LOG.info( + "[Agentic RAG] Starting research — %d message(s), last role=%s, content_len=%d", + len(messages), + messages[-1].get("role", "") if messages else "?", + len(messages[-1].get("content", "")) if messages else 0, + ) token_queue: asyncio.Queue = asyncio.Queue() graph = build_agentic_graph(tools, token_queue, gen_conf=gen_conf) diff --git a/rag/advanced_rag/harness/config.py b/rag/advanced_rag/harness/config.py index 0266ec4716..0c996d6469 100644 --- a/rag/advanced_rag/harness/config.py +++ b/rag/advanced_rag/harness/config.py @@ -53,6 +53,7 @@ THINKING_MODES: dict[str, ExecutionStrategy] = { "hybrid_search", "web_search", "catalog_navigate", + "dataset_navigate", "graph_explore", "inspector_open_context", "inspector_compare", @@ -79,6 +80,7 @@ THINKING_MODES: dict[str, ExecutionStrategy] = { "web_search", "structured_query", "catalog_navigate", + "dataset_navigate", "mindmap_navigate", "graph_explore", "wiki_query", diff --git a/rag/advanced_rag/harness/orchestrator/agentic.py b/rag/advanced_rag/harness/orchestrator/agentic.py index 9ed7a6463d..19bb9424f0 100644 --- a/rag/advanced_rag/harness/orchestrator/agentic.py +++ b/rag/advanced_rag/harness/orchestrator/agentic.py @@ -91,7 +91,7 @@ async def agentic_research(state: dict, tools) -> dict: description=dc, ) ) - _LOG.info("[Agentic research] Found a new angle worth researching: \"%s\"", dc) + _LOG.info('[Agentic research] Found a new angle worth researching: "%s"', dc) # ── Step B: Sufficiency Check ── all_chunks = {i: c for i, c in enumerate(tools.kbinfos.get("chunks", []))} @@ -140,7 +140,7 @@ async def _run_claim_research( mode, compilation_map: dict, ) -> dict: - _LOG.info("[Agentic research] Researching: \"%s\"", _snip(claim.description)) + _LOG.info('[Agentic research] Researching: "%s"', _snip(claim.description)) try: result = await asyncio.wait_for( research_agent_loop(claim, tools, pipeline, ctx, mode, compilation_map), @@ -150,7 +150,7 @@ async def _run_claim_research( raise except asyncio.TimeoutError: _LOG.warning( - "[Agentic research] Gave up on \"%s\" — it took longer than %ss.", + '[Agentic research] Gave up on "%s" — it took longer than %ss.', _snip(claim.description), CLAIM_RESEARCH_TIMEOUT_SECONDS, ) @@ -163,7 +163,7 @@ async def _run_claim_research( "discovered_claims": [], } except Exception: - _LOG.exception("[Agentic research] Hit an error while researching \"%s\".", _snip(claim.description)) + _LOG.exception('[Agentic research] Hit an error while researching "%s".', _snip(claim.description)) return { "report": "", "is_verified": False, @@ -174,7 +174,7 @@ async def _run_claim_research( } _LOG.info( - "[Agentic research] Finished \"%s\" — %s, backed by %d passage(s) (confidence %.0f%%)%s.", + '[Agentic research] Finished "%s" — %s, backed by %d passage(s) (confidence %.0f%%)%s.', _snip(claim.description), "answered" if result.get("is_verified") else "still unanswered", len(result.get("evidence_ids") or []), @@ -233,6 +233,91 @@ async def _get_compilation_map(tools) -> dict[str, set[str]]: comps.add("mindmap") if parser_config.get("page_index"): comps.add("page_index") + await _add_template_group_compilations(comps, parser_config, getattr(kb, "tenant_id", "")) + if await _has_dataset_nav_rows(getattr(kb, "tenant_id", ""), getattr(kb, "id", "")): + comps.add("tree") if comps: result[kb.id] = comps return result + + +async def _has_dataset_nav_rows(tenant_id: str, kb_id: str) -> bool: + if not tenant_id or not kb_id: + return False + try: + from common import settings + from common.doc_store.doc_store_base import OrderByExpr + from common.misc_utils import thread_pool_exec + from rag.nlp import search + + index_name = search.index_name(tenant_id) + if not settings.docStoreConn.index_exist(index_name, kb_id): + return False + res = await thread_pool_exec( + settings.docStoreConn.search, + ["id"], + [], + {"compile_kwd": ["dataset_nav"]}, + [], + OrderByExpr(), + 0, + 1, + index_name, + [kb_id], + ) + return bool(settings.docStoreConn.get_total(res)) + except Exception: + _LOG.exception("[agentic] dataset-nav existence check failed for kb=%s", kb_id) + return False + + +async def _add_template_group_compilations(comps: set[str], parser_config: dict, tenant_id: str) -> None: + """Infer available compilation kinds from selected template groups.""" + if not tenant_id: + return + try: + from common.misc_utils import thread_pool_exec + from api.db.services.compilation_template_group_service import CompilationTemplateGroupService + from rag.svr.task_executor_refactor.chunk_post_processor import ( + _parser_config_compilation_template_group_ids, + ) + except Exception: + _LOG.exception("[agentic] compilation-map helper import failed") + return + + try: + group_ids = _parser_config_compilation_template_group_ids(parser_config) + except Exception: + _LOG.exception("[agentic] compilation template group id resolution failed") + return + + for group_id in group_ids: + try: + group = await thread_pool_exec(CompilationTemplateGroupService.get_saved, group_id, tenant_id) + except Exception: + _LOG.exception("[agentic] compilation template group read failed id=%s", group_id) + continue + for template in (group or {}).get("templates") or []: + config = template.get("config") or {} + raw_kind = (config.get("kind") if isinstance(config, dict) else "") or template.get("kind") or "" + raw_norm = raw_kind.strip().lower().replace("-", "_") if isinstance(raw_kind, str) else "" + kind = _compilation_kind_for_agentic_map(raw_kind) + if raw_norm == "knowledge_graph": + comps.add("knowledge_graph") + if kind == "tree": + comps.add("tree") + elif kind in {"timeline", "page_index", "pageindex"}: + comps.add("page_index") + elif kind in {"mindmap", "mind_map"}: + comps.add("mindmap") + elif kind == "artifacts": + comps.add("wiki") + + +def _compilation_kind_for_agentic_map(kind) -> str: + if not isinstance(kind, str): + return "" + normalized = kind.strip().lower().replace("-", "_") + if normalized in {"pageindex", "page_index"}: + return "timeline" + return normalized diff --git a/rag/advanced_rag/harness/tools/__init__.py b/rag/advanced_rag/harness/tools/__init__.py index de58a9c1af..63954f575e 100644 --- a/rag/advanced_rag/harness/tools/__init__.py +++ b/rag/advanced_rag/harness/tools/__init__.py @@ -14,10 +14,9 @@ register_tool("web_search", _search_schema("web_search", "Internet search"), web register_tool("structured_query", _search_schema("structured_query", "SQL search"), structured_query) # Navigation tools (require compilation) -from rag.advanced_rag.harness.tools.navigation import catalog_navigate, mindmap_navigate +from rag.advanced_rag.harness.tools.navigation import catalog_navigate, dataset_navigate, mindmap_navigate -# catalog_navigate covers both the tree/TOC outline and the page index — it reads -# whichever compiled catalog the doc has, so it is offered when either exists. +# catalog_navigate covers both the tree/TOC outline and the page index. register_tool( "catalog_navigate", _navigate_schema("catalog_navigate", "Answer from the document's compiled catalog (table of contents / page index)"), @@ -26,6 +25,13 @@ register_tool( compilation_type=("toc", "page_index"), ) register_tool("mindmap_navigate", _navigate_schema("mindmap_navigate", "Navigate by mindmap"), mindmap_navigate, requires_compilation=True, compilation_type="mindmap") +register_tool( + "dataset_navigate", + _navigate_schema("dataset_navigate", "Find the most relevant documents via the dataset map, then search within them"), + dataset_navigate, + requires_compilation=True, + compilation_type="tree", +) # Exploration tools (require compilation) from rag.advanced_rag.harness.tools.exploration import graph_explore, wiki_query diff --git a/rag/advanced_rag/harness/tools/gating.py b/rag/advanced_rag/harness/tools/gating.py index f0c231fe53..fb208e44f4 100644 --- a/rag/advanced_rag/harness/tools/gating.py +++ b/rag/advanced_rag/harness/tools/gating.py @@ -10,6 +10,7 @@ SEARCH_PHASES = { "locate": { "goal": "Locate documents or regions that may contain the answer.", "tools_priority": [ + "dataset_navigate", "catalog_navigate", "mindmap_navigate", "wiki_query", @@ -65,7 +66,8 @@ def compilation_available(tool_name: str, compilation_map: dict) -> bool: comp_type = tool["compilation_type"] if not compilation_map: return False - return any(comp_type in comps for comps in compilation_map.values()) + comp_types = set(comp_type) if isinstance(comp_type, (list, tuple, set)) else {comp_type} + return any(bool(comp_types & set(comps)) for comps in compilation_map.values()) def tool_fits_context(tool_name: str, context: OrchestratorContext) -> bool: @@ -74,6 +76,8 @@ def tool_fits_context(tool_name: str, context: OrchestratorContext) -> bool: return False if tool_name == "catalog_navigate" and not context.current_claim: return False + if tool_name == "dataset_navigate" and not context.current_claim: + return False if tool_name == "graph_explore" and not context.last_entity: return False if tool_name == "mindmap_navigate" and not context.current_claim: diff --git a/rag/advanced_rag/harness/tools/navigation.py b/rag/advanced_rag/harness/tools/navigation.py index 345ae64e24..989c0be680 100644 --- a/rag/advanced_rag/harness/tools/navigation.py +++ b/rag/advanced_rag/harness/tools/navigation.py @@ -374,6 +374,82 @@ async def mindmap_navigate(tools, topic: str, keywords: str = "", doc_scope: lis ) +# ── Dataset navigation (document router) ──────────────────────────────────── + +_NAV_MAX_DOCS = 8 # documents the nav tree routes a query to +_NAV_MAX_HITS_PER_KB = 8 + + +async def dataset_navigate(tools, topic: str, keywords: str = "", doc_scope: list[str] | None = None) -> dict: + """Route the query to the most relevant documents via the dataset nav tree, + then search within only those documents. + + The nav tree is a KB-level RAPTOR-style summary of every document; searching + it narrows the corpus to the handful of documents worth reading before the + real chunk retrieval runs (coarse-to-fine). Falls back to a plain hybrid + search when the KB has no nav tree or nothing routes. + + :returns: ``{"answer": "", "chunks": [...], "doc_aggs": [...]}`` + """ + from rag.advanced_rag.harness.tools.search import hybrid_search + + query = topic + _LOG.info(f'[Dataset navigation] Finding the most relevant documents for "{query}" (keywords: {keywords})') + + # Caller already narrowed the corpus — just retrieve within it. + if doc_scope: + return await hybrid_search(tools, query=query, keywords=keywords, doc_scope=doc_scope) + + from rag.advanced_rag.knowlege_compile.dataset_nav import search_dataset_nav + + candidates: list[tuple[float, str]] = [] + nav_entities: list[dict] = [] + seen: set[str] = set() + seen_nodes: set[str] = set() + for kb in getattr(tools, "kbs", []) or []: + try: + hits = await search_dataset_nav( + kb.tenant_id, + kb.id, + query, + embd_mdl=getattr(tools, "embed_mdl", None), + top_k=_NAV_MAX_HITS_PER_KB, + ) + except Exception: + _LOG.exception("[Dataset navigation] nav-tree search failed for kb=%s", kb.id) + continue + for h in hits: + score = float(h.get("score") or 0.0) + node_name = str(h.get("name") or "") + if node_name and node_name not in seen_nodes: + seen_nodes.add(node_name) + nav_entities.append( + { + "name": node_name, + "type": h.get("type") or "dataset_nav", + "discription": h.get("description") or "", + } + ) + for did in h.get("doc_ids") or []: + if did and did not in seen: + seen.add(did) + candidates.append((score, did)) + + candidates.sort(key=lambda item: item[0], reverse=True) + routed = [did for _, did in candidates[:_NAV_MAX_DOCS]] + if not routed: + _LOG.info("[Dataset navigation] No dataset map here — falling back to a normal search.") + return await hybrid_search(tools, query=query, keywords=keywords) + + _LOG.info("[Dataset navigation] Routed to %d document(s); searching within them.", len(routed)) + answer = "" + if nav_entities: + answer, _ = await _ask_structure(tools, query, nav_entities, [], "dataset map", "Dataset navigation") + result = await hybrid_search(tools, query=query, keywords=keywords, doc_scope=routed) + result["answer"] = answer + return result + + # ── Knowledge-graph exploration ───────────────────────────────────────────── # # Unlike catalog/mindmap (which read the merged "graph" JSON of one doc), the KG diff --git a/rag/advanced_rag/harness/tools/registry.py b/rag/advanced_rag/harness/tools/registry.py index 74bcb2db73..aea623b9af 100644 --- a/rag/advanced_rag/harness/tools/registry.py +++ b/rag/advanced_rag/harness/tools/registry.py @@ -11,7 +11,7 @@ TOOL_REGISTRY: dict[str, dict[str, Any]] = {} # async def fn(tools, **kwargs) -> dict # {"chunks": [...], ...} -def register_tool(name: str, schema: dict, fn: callable, requires_compilation: bool = False, compilation_type: str | None = None, processing_time: str = "fast") -> None: +def register_tool(name: str, schema: dict, fn: callable, requires_compilation: bool = False, compilation_type: str | tuple[str, ...] | None = None, processing_time: str = "fast") -> None: TOOL_REGISTRY[name] = { "name": name, "function_schema": schema, diff --git a/rag/advanced_rag/knowlege_compile/dataset_nav.py b/rag/advanced_rag/knowlege_compile/dataset_nav.py index 9696cb70f7..54b249fa9e 100644 --- a/rag/advanced_rag/knowlege_compile/dataset_nav.py +++ b/rag/advanced_rag/knowlege_compile/dataset_nav.py @@ -29,6 +29,7 @@ from __future__ import annotations import asyncio import json import logging +import re from typing import Any import xxhash @@ -181,6 +182,8 @@ async def _store_knn( fields = [ "content_with_weight", "name", + "doc_id", + "type_kwd", "parent_kwd", "depth_int", "doc_count_int", @@ -926,6 +929,118 @@ async def _maybe_split_cluster( await _store_upsert(tenant_id, kb_id, row) +async def search_dataset_nav( + tenant_id: str, + kb_id: str, + query: str, + embd_mdl=None, + top_k: int = 8, +) -> 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:: + + {"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. + """ + query = (query or "").strip() + if not query: + return [] + + condition = {"compile_kwd": [_COMPILE_KWD]} + rows_with_scores: list[tuple[dict, float]] = [] + + if embd_mdl is not None: + try: + vec = await _embed(embd_mdl, query) + except Exception: + logging.exception("search_dataset_nav: embed failed for kb=%s", kb_id) + vec = [] + if vec: + try: + rows = await _store_knn(tenant_id, kb_id, vec, len(vec), condition, top_k=top_k) + vf = _vec_field(len(vec)) + rows_with_scores = [(r, _cosine_sim(vec, r.get(vf) or [])) for r in rows] + 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) + + # 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] + + out: list[dict] = [] + for r, score in rows_with_scores[:top_k]: + try: + payload = json.loads(r.get("content_with_weight") or "{}") + except Exception: + payload = {} + typ = payload.get("type") or r.get("type_kwd") or ("nav_cluster" if r.get("doc_ids_kwd") else "nav_doc") + name = r.get("name") or "" + if typ == "nav_cluster": + doc_id = None + doc_ids = _as_str_list(r.get("doc_ids_kwd")) + else: + # Leaf: ``name`` == the document id (see ``_make_nav_doc_row``). + doc_id = r.get("doc_id") or name + doc_ids = [doc_id] if doc_id else [] + out.append( + { + "type": typ, + "doc_id": doc_id, + "doc_ids": doc_ids, + "name": name, + "description": payload.get("description") or "", + "doc_count": int(r.get("doc_count_int") or len(doc_ids) or 0), + "score": float(score or 0.0), + } + ) + return out + + +def _as_str_list(value) -> list[str]: + if isinstance(value, list): + return [str(v) for v in value if v] + if isinstance(value, str) and value: + return [value] + return [] + + +def _nav_text_score(query: str, row: dict) -> float: + try: + payload = json.loads(row.get("content_with_weight") or "{}") + except Exception: + payload = {} + haystack = " ".join( + str(x or "") + for x in ( + row.get("name"), + payload.get("description"), + ) + ).lower() + q_terms = set(re.findall(r"[\w]+", query.lower())) + if not q_terms: + return 0.0 + hits = sum(1 for term in q_terms if term in haystack) + return hits / max(len(q_terms), 1) + + def remove_dataset_nav_doc_sync( tenant_id: str, kb_id: str,