Let agentic rag search honor metadata filter (#17731)

rag/advanced_rag/agentic_rag.py

@@ -77,6 +77,7 @@ class RAGTools:
         kbs: list[Knowledgebase] | None = None,
         tav: Tavily | None = None,
         meta_data_filter: dict | None = None,
+        doc_scope: List[str] | None = None,
         user_defined_prompts: dict | None = None,
         do_refer: bool | None = True,
         thinking_mode: str = "medium",
@@ -107,6 +108,7 @@ class RAGTools:
 
         self.tav = tav
         self.meta_data_filter = meta_data_filter
+        self.doc_scope = list(dict.fromkeys(doc_scope)) if doc_scope is not None else None
         self.user_defined_prompts = user_defined_prompts or {}
         self.kbinfos = {"chunks": [], "doc_aggs": []}
         self.do_refer = do_refer
@@ -140,6 +142,14 @@ class RAGTools:
     def has_llm(self) -> bool:
         return self.chat_mdl is not None
 
+    def scoped_doc_ids(self, doc_scope: List[str] | None = None) -> List[str] | None:
+        if self.doc_scope is None:
+            return doc_scope
+        if not doc_scope:
+            return list(self.doc_scope)
+        allowed = set(self.doc_scope)
+        return [doc_id for doc_id in doc_scope if doc_id in allowed]
+
     async def _fit_messages(self, system: str, user: str) -> list:
         """Fit system+user messages into the model's context window."""
         from rag.prompts.generator import form_message, message_fit_in
@@ -369,6 +379,9 @@ class RAGTools:
             keywords = ",".join(keywords)
         logging.info(f"@retrieve: {question}@{keywords}")
 
+        doc_scope = self.scoped_doc_ids(doc_scope)
+        if doc_scope == ["-999"]:
+            return {"chunks": [], "doc_aggs": []}
         if doc_scope:
             candidates = [d for d in doc_scope if isinstance(d, str)]
             known = await thread_pool_exec(self._filter_known_doc_ids, candidates)
@@ -376,6 +389,8 @@ class RAGTools:
             if valid:
                 doc_scope = valid
             else:
+                if self.doc_scope is not None:
+                    return {"chunks": [], "doc_aggs": []}
                 if candidates:
                     logging.warning("retrieve: every supplied doc ID was unknown; falling back to unfiltered retrieval")
                 doc_scope = None
@@ -434,7 +449,7 @@ class RAGTools:
         sql_kb_ids = [kb.id for kb in self.sql_kbs]
         tenant_id = self.sql_kbs[0].tenant_id
         try:
-            ans = await use_sql(question, self.field_map, tenant_id, self.chat_mdl, quota=True, kb_ids=sql_kb_ids)
+            ans = await use_sql(question, self.field_map, tenant_id, self.chat_mdl, quota=True, kb_ids=sql_kb_ids, doc_ids=self.scoped_doc_ids())
         except Exception as e:
             logging.exception(f"structured_retrieve: use_sql failed: {e}")
             return {"answer": "", "chunks": [], "doc_aggs": []}
@@ -491,6 +506,8 @@ class RAGTools:
         """Fetch a whole document's chunks in reading order (raw kbinfos)."""
         if not self.kb_ids:
             return {"chunks": [], "doc_aggs": []}
+        if self.doc_scope is not None and doc_id not in self.doc_scope:
+            return {"chunks": [], "doc_aggs": []}
         resolved = await thread_pool_exec(self._resolve_doc_tenant, doc_id)
         if resolved is None:
             logging.warning(f"fetch_full_document: doc_id {doc_id!r} not in any bound KB — refusing to fetch")

rag/advanced_rag/harness/pipeline.py

@@ -13,7 +13,7 @@ _LOG = logging.getLogger(__name__)
 # (``dataset_navigation_by_tree``) has produced a relevant-document set, these
 # inherit it as their ``doc_scope`` unless the caller passed one explicitly, so
 # a follow-up search stays within the routed docs instead of re-scanning the KB.
-_DOC_SCOPE_CONSUMERS = {"ontology_navigate", "mindmap_navigate", "graph_explore", "hybrid_search"}
+_DOC_SCOPE_CONSUMERS = {"ontology_navigate", "mindmap_navigate", "graph_explore", "hybrid_search", "vector_search", "bm25_search", "structured_query", "dataset_navigation_by_tree"}
 
 
 class Pipeline:
@@ -30,7 +30,7 @@ class Pipeline:
         self.compilation_map = compilation_map or {}
         self.trace: list[dict] = []
         # Latest relevant-document set produced by a routing tool this run.
-        self._routed_docs: list[str] = []
+        self._routed_docs: list[str] = list(getattr(rag_tools, "doc_scope", None) or [])
 
     async def execute(self, tool_name: str, **kwargs) -> ToolResult:
         """Execute a registered tool by name."""
@@ -58,7 +58,10 @@ class Pipeline:
             # document IDs; remember them so the scope-consuming tools above can
             # inherit them on later turns.
             if result.docs:
-                self._routed_docs = list(result.docs)
+                if hasattr(self.tools, "scoped_doc_ids"):
+                    self._routed_docs = self.tools.scoped_doc_ids(list(result.docs)) or []
+                else:
+                    self._routed_docs = list(result.docs)
             # Feed the shared citation pool: agent searches go through the
             # pipeline, so without this their evidence never reaches kbinfos and
             # the final answer has nothing to cite.

rag/advanced_rag/harness/tools/navigation.py

@@ -314,6 +314,8 @@ async def ontology_navigate(tools, topic: str, keywords: str = "", doc_scope: li
 
     :returns: ``{"answer": "", "chunks": [...], "doc_aggs": [...]}``
     """
+    if hasattr(tools, "scoped_doc_ids"):
+        doc_scope = tools.scoped_doc_ids(doc_scope)
     if not doc_scope:
         doc_scope = []
     _LOG.info(f'[Ontology navigation] Looking through the document catalog for "{topic}" (keywords: {keywords}) in doc: {len(doc_scope)}')
@@ -332,6 +334,8 @@ async def mindmap_navigate(tools, topic: str, keywords: str = "", doc_scope: lis
 
     :returns: ``{"answer": "", "chunks": [...], "doc_aggs": [...]}``
     """
+    if hasattr(tools, "scoped_doc_ids"):
+        doc_scope = tools.scoped_doc_ids(doc_scope)
     if not doc_scope:
         doc_scope = []
     _LOG.info(f'[Mindmap navigation] Following the concept mindmap for "{topic}" (keywords: {keywords}) in doc: {len(doc_scope)}')
@@ -418,7 +422,7 @@ async def _ask_nav_select(tools, query: str, items: list[dict], noun: str, max_i
     return out
 
 
-async def _collect_nav_leaves(dataset_api_service, clusters: list[dict]) -> list[dict]:
+async def _collect_nav_leaves(dataset_api_service, clusters: list[dict], doc_scope: list[str] | None = None) -> list[dict]:
     """BFS from the selected clusters down to their document leaves.
 
     Each cluster carries ``name`` + ``kb``. A node's children are either document
@@ -429,6 +433,7 @@ async def _collect_nav_leaves(dataset_api_service, clusters: list[dict]) -> list
     seen_docs: set[str] = set()
     seen_nodes: set[tuple] = set()
     frontier: list[tuple] = [(c["kb"], c["name"], 0) for c in clusters if c.get("name")]
+    allowed_docs = set(doc_scope or [])
 
     while frontier and len(leaves) < _NAV_TREE_MAX_LEAVES:
         kb, name, depth = frontier.pop(0)
@@ -446,7 +451,7 @@ async def _collect_nav_leaves(dataset_api_service, clusters: list[dict]) -> list
         for item in data.get("items") or []:
             if item.get("type") == "doc":
                 did = str(item.get("doc_id") or "").strip()
-                if did and did not in seen_docs:
+                if did and (not allowed_docs or did in allowed_docs) and did not in seen_docs:
                     seen_docs.add(did)
                     leaves.append({**item, "kb": kb})
                     if len(leaves) >= _NAV_TREE_MAX_LEAVES:
@@ -481,6 +486,8 @@ async def dataset_navigation_by_tree(tools, topic: str, keywords: str = "", doc_
     query = " ".join(part for part in ((topic or "").strip(), (keywords or "").strip()) if part).strip()
     if not query:
         return []
+    if hasattr(tools, "scoped_doc_ids"):
+        doc_scope = tools.scoped_doc_ids(doc_scope)
 
     _LOG.info('[Dataset navigation] Walking the dataset tree for "%s"', query)
 
@@ -514,7 +521,7 @@ async def dataset_navigation_by_tree(tools, topic: str, keywords: str = "", doc_
     _LOG.info("[Dataset navigation] %d/%d cluster(s) selected.", len(selected_clusters), len(clusters))
 
     # 3. Descend the selected clusters to their document leaves.
-    leaves = await _collect_nav_leaves(dataset_api_service, selected_clusters)
+    leaves = await _collect_nav_leaves(dataset_api_service, selected_clusters, doc_scope)
     if not leaves:
         _LOG.info("[Dataset navigation] no leaf under selected cluster %s.", _nav_cluster_names(selected_clusters))
         return []
@@ -566,6 +573,8 @@ async def _kg_scopes(tools, doc_scope: list[str] | None = None):
     """
     from common.misc_utils import thread_pool_exec
 
+    if hasattr(tools, "scoped_doc_ids"):
+        doc_scope = tools.scoped_doc_ids(doc_scope)
     if doc_scope:
         by_kb: dict[tuple, list[str]] = {}
         for doc_id in doc_scope:
@@ -759,6 +768,8 @@ async def graph_explore(tools, query: str, keywords: str = "", doc_scope: list[s
     from rag.advanced_rag.harness.tools.search import _narrow_by_keywords
 
     _empty = {"answer": "", "chunks": [], "doc_aggs": []}
+    if hasattr(tools, "scoped_doc_ids"):
+        doc_scope = tools.scoped_doc_ids(doc_scope)
     _LOG.info(f'[Graph exploration] Exploring the knowledge graph for "{query}" (keywords: {keywords})')
 
     scopes = await _kg_scopes(tools, doc_scope)

rag/advanced_rag/harness/tools/search.py

@@ -184,6 +184,8 @@ async def hybrid_search(tools, query: str, kb_ids: list[str] | None = None, top_
     if not tools.kb_ids and not kb_ids:
         return {"chunks": [], "doc_aggs": []}
     target_ids = kb_ids or tools.kb_ids
+    if hasattr(tools, "scoped_doc_ids"):
+        doc_scope = tools.scoped_doc_ids(doc_scope)
     _LOG.info(f'[Hybrid search] Searching the knowledge base for "{query}" (keywords: {keywords})')
 
     # Query expansion: append the formalized-question keywords + close synonyms
@@ -223,13 +225,13 @@ async def hybrid_search(tools, query: str, kb_ids: list[str] | None = None, top_
         _LOG.info(f"[Hybrid search] Kept {len(kbinfos['chunks'])} of {length} passage(s) that actually mention the keywords.")
     if use_compiled and kbinfos.get("chunks"):
         _LOG.info("[Hybrid search] Compiled expansion enabled — enriching with page_index/tree/KG navigation.")
-        await _expand_with_compiled(tools, query, keywords, kbinfos)
+        await _expand_with_compiled(tools, query, keywords, kbinfos, doc_scope)
     if cache is not None:
         cache[cache_key] = kbinfos
     return kbinfos
 
 
-async def vector_search(tools, query: str, kb_ids: list[str] | None = None, top_n: int = 12, keywords: str = "") -> dict:
+async def vector_search(tools, query: str, kb_ids: list[str] | None = None, top_n: int = 12, keywords: str = "", doc_scope: list[str] | None = None) -> dict:
     if not tools.embed_mdl:
         _LOG.warning("vector_search: no embed_mdl available")
         return {"chunks": [], "doc_aggs": []}
@@ -237,6 +239,8 @@ async def vector_search(tools, query: str, kb_ids: list[str] | None = None, top_
     _LOG.info(f'[Vector search] Searching by meaning for "{query}" (keywords: {keywords})')
     effective_query = f"{query} {keywords}".strip() if keywords else query
     target_ids = kb_ids or tools.kb_ids
+    if hasattr(tools, "scoped_doc_ids"):
+        doc_scope = tools.scoped_doc_ids(doc_scope)
     kbinfos = await settings.retriever.retrieval(
         effective_query,
         tools.embed_mdl,
@@ -248,6 +252,7 @@ async def vector_search(tools, query: str, kb_ids: list[str] | None = None, top_
         vector_similarity_weight=1.0,
         aggs=False,
         highlight=False,
+        doc_ids=doc_scope,
     )
     kbinfos = _normalize(kbinfos, tools.tenant_ids)
     if keywords:
@@ -257,10 +262,12 @@ async def vector_search(tools, query: str, kb_ids: list[str] | None = None, top_
     return kbinfos
 
 
-async def bm25_search(tools, query: str, kb_ids: list[str] | None = None, top_n: int = 12, keywords: str = "") -> dict:
+async def bm25_search(tools, query: str, kb_ids: list[str] | None = None, top_n: int = 12, keywords: str = "", doc_scope: list[str] | None = None) -> dict:
     _LOG.info(f'[BM25 search] Searching by keyword for "{query}" (keywords: {keywords})')
     target_ids = kb_ids or tools.kb_ids
     effective_query = f"{query} {keywords}".strip() if keywords else query
+    if hasattr(tools, "scoped_doc_ids"):
+        doc_scope = tools.scoped_doc_ids(doc_scope)
     kbinfos = await settings.retriever.retrieval(
         effective_query,
         None,
@@ -272,6 +279,7 @@ async def bm25_search(tools, query: str, kb_ids: list[str] | None = None, top_n:
         vector_similarity_weight=0,
         aggs=False,
         highlight=False,
+        doc_ids=doc_scope,
     )
     kbinfos = _normalize(kbinfos, tools.tenant_ids)
     if keywords:
@@ -284,7 +292,7 @@ async def bm25_search(tools, query: str, kb_ids: list[str] | None = None, top_n:
 # ─── Compiled product expansion (zero-LLM, used by hybrid_search with use_compiled=True) ───
 
 
-async def _expand_with_compiled(tools, query: str, keywords: str, kbinfos: dict) -> None:
+async def _expand_with_compiled(tools, query: str, keywords: str, kbinfos: dict, doc_scope: list[str] | None = None) -> None:
     """Zero-LLM compiled-product expansion: page_index → tree → KG.
 
     For each bound KB, searches compiled entity rows matching the query,
@@ -294,7 +302,7 @@ async def _expand_with_compiled(tools, query: str, keywords: str, kbinfos: dict)
     before = len(kbinfos.get("chunks", []))
     seen_ids = {c.get("chunk_id") or c.get("id") for c in kbinfos.get("chunks", [])}
 
-    scopes = await _kg_scopes(tools)
+    scopes = await _kg_scopes(tools, doc_scope)
     if not scopes:
         return
 
@@ -764,7 +772,7 @@ async def web_search(tools, query: str, keywords: str = "") -> dict:
         return {"chunks": [], "doc_aggs": []}
 
 
-async def structured_query(tools, query: str, keywords: str = "", kb_ids: list[str] | None = None) -> dict:
+async def structured_query(tools, query: str, keywords: str = "", kb_ids: list[str] | None = None, doc_scope: list[str] | None = None) -> dict:
     """Answer from the structured (tabular) KBs by translating the query to SQL.
 
     ``keywords`` is accepted for schema conformance but deliberately unused: the
@@ -775,12 +783,14 @@ async def structured_query(tools, query: str, keywords: str = "", kb_ids: list[s
     sql_kbs = [kb for kb in tools.sql_kbs if kb_ids is None or kb.id in kb_ids]
     if not sql_kbs:
         return {"answer": "", "chunks": [], "doc_aggs": []}
+    if hasattr(tools, "scoped_doc_ids"):
+        doc_scope = tools.scoped_doc_ids(doc_scope)
     from api.db.services.dialog_service import use_sql
 
     tenant_id = sql_kbs[0].tenant_id
     sql_kb_ids = [kb.id for kb in sql_kbs]
     try:
-        ans = await use_sql(query, tools.field_map, tenant_id, tools.chat_mdl, quota=True, kb_ids=sql_kb_ids)
+        ans = await use_sql(query, tools.field_map, tenant_id, tools.chat_mdl, quota=True, kb_ids=sql_kb_ids, doc_ids=doc_scope)
     except Exception:
         _LOG.exception("structured_query failed")
         return {"answer": "", "chunks": [], "doc_aggs": []}