From fac40e510323dc63849e5a087fd1e27fa184df1e Mon Sep 17 00:00:00 2001 From: Kevin Hu Date: Tue, 4 Aug 2026 18:02:13 +0800 Subject: [PATCH] Refactor: Make wiki and web searchable. (#17789) ### Summary Refine wiki and web searchable. Closes #17638 --- rag/advanced_rag/harness/agent.py | 1 + rag/advanced_rag/harness/tools/exploration.py | 113 ++++++++++++++++-- rag/advanced_rag/harness/tools/gating.py | 17 ++- rag/advanced_rag/harness/tools/search.py | 3 + rag/advanced_rag/knowlege_compile/wiki.py | 44 ++++++- rag/llm/chat_model.py | 79 +++++++++++- rag/nlp/search.py | 3 + 7 files changed, 237 insertions(+), 23 deletions(-) diff --git a/rag/advanced_rag/harness/agent.py b/rag/advanced_rag/harness/agent.py index c12a6ab286..d1e5a3f377 100644 --- a/rag/advanced_rag/harness/agent.py +++ b/rag/advanced_rag/harness/agent.py @@ -155,6 +155,7 @@ async def research_agent_loop( compilation_map=compilation_map, context=context, has_routed_scope=bool(getattr(pipeline, "_routed_docs", None)), + web_enabled=bool(getattr(tools, "has_web", lambda: False)()), ) # Clone so binding tools never leaks onto the shared chat model. diff --git a/rag/advanced_rag/harness/tools/exploration.py b/rag/advanced_rag/harness/tools/exploration.py index f44856c1e1..3cc75763e8 100644 --- a/rag/advanced_rag/harness/tools/exploration.py +++ b/rag/advanced_rag/harness/tools/exploration.py @@ -3,13 +3,15 @@ ``graph_explore`` lives in :mod:`navigation` (it shares the compiled-structure machinery) and is re-exported here so the tool registry keeps one import point. -``wiki_query`` retrieves through ``hybrid_search`` and takes the same -``keywords`` the other search tools do — the keywords drive query expansion and -the keyword-sentence narrowing. Parameter names must match the registered -``_search_schema`` (``query`` + ``keywords``), otherwise every LLM tool call -fails with a TypeError. +``wiki_query`` runs a hybrid (BM25 + dense) search over the searchable wiki draft +rows written by ``_wiki_persist_draft`` (``compile_kwd="wiki_page_draft"``) and +returns each page's markdown as a chunk. It takes the same ``keywords`` the other +search tools do — the keywords drive the keyword-sentence narrowing. Parameter +names must match the registered ``_search_schema`` (``query`` + ``keywords``), +otherwise every LLM tool call fails with a TypeError. """ +import json import logging # graph_explore is implemented alongside catalog/mindmap navigation because it @@ -18,15 +20,102 @@ from rag.advanced_rag.harness.tools.navigation import graph_explore # noqa: F40 _LOG = logging.getLogger(__name__) +# compile_kwd of the searchable wiki draft rows (see _wiki_persist_draft). +_WIKI_DRAFT_COMPILE_KWD = "wiki_page_draft" +_WIKI_QUERY_TOP_N = 12 + async def wiki_query(tools, query: str, keywords: str = "") -> dict: - """Query compiled wiki knowledge. + """Search the compiled wiki. - This is currently a placeholder. The final implementation should call the - compiled wiki store. + Hybrid (BM25 over ``title_tks`` / ``content_ltks`` / ``content_sm_ltks`` + + dense over ``q__vec``) search across each bound KB's ``wiki_page_draft`` + rows. The page markdown is parsed out of each row's ``content_with_weight`` + (which stays the page JSON) and returned as chunks, narrowed by ``keywords``. + + :returns: ``{"answer": "", "chunks": [...], "doc_aggs": [...]}`` """ - _LOG.info(f'[Wiki lookup] Looking up the compiled wiki for "{query}" (keywords: {keywords})') - # TODO: implement actual wiki lookup - from rag.advanced_rag.harness.tools.search import hybrid_search + from common import settings + from common.doc_store.doc_store_base import FusionExpr, OrderByExpr + from common.misc_utils import thread_pool_exec + from rag.nlp import search as _rag_search + from rag.advanced_rag.harness.tools.search import _narrow_by_keywords - return await hybrid_search(tools, query=query, keywords=keywords) + _LOG.info(f'[Wiki lookup] Searching the compiled wiki for "{query}" (keywords: {keywords})') + + kbs = getattr(tools, "kbs", []) or [] + text = f"{query} {keywords}".strip() + if not kbs or not text: + return {"answer": "", "chunks": [], "doc_aggs": []} + + fields = ["content_with_weight", "docnm_kwd", "title_kwd", "wiki_slug_kwd", "source_doc_ids", "doc_id"] + qryr = settings.retriever.qryr + chunks: list[dict] = [] + + for kb in kbs: + kb_id = kb.id + tenant_id = kb.tenant_id + index = _rag_search.index_name(tenant_id) + try: + # BM25 over the standard tokenized fields, fused with dense when an + # embedder is available — mirrors the retriever's own hybrid search. + match_text, _ = qryr.question(text, min_match=0.3) + exprs = [match_text] + if getattr(tools, "embed_mdl", None): + try: + match_dense = await settings.retriever.get_vector(text, tools.embed_mdl, _WIKI_QUERY_TOP_N, 0.1) + exprs = [match_text, match_dense, FusionExpr("weighted_sum", _WIKI_QUERY_TOP_N, {"weights": "0.001, 1"})] + except Exception: + _LOG.exception("[Wiki lookup] dense expr build failed; BM25 only") + res = await thread_pool_exec( + settings.docStoreConn.search, + fields, + [], + {"compile_kwd": [_WIKI_DRAFT_COMPILE_KWD]}, + exprs, + OrderByExpr(), + 0, + _WIKI_QUERY_TOP_N, + index, + [kb_id], + ) + rows = settings.docStoreConn.get_fields(res, fields) or {} + except Exception: + _LOG.exception("[Wiki lookup] search failed for kb=%s", kb_id) + continue + + for cid, row in rows.items(): + try: + page = json.loads(row.get("content_with_weight") or "{}") + except Exception: + page = {} + if not isinstance(page, dict): + page = {} + content = page.get("content_md_rendered") or page.get("content_md") or page.get("content_md_raw") or "" + if not content: + continue + title = row.get("docnm_kwd") or page.get("title") or row.get("title_kwd") or "" + slug = row.get("wiki_slug_kwd") or page.get("slug") or "" + chunks.append( + { + "chunk_id": cid, + "content_with_weight": content, + "docnm_kwd": title, + "doc_id": slug or row.get("doc_id") or kb_id, + "wiki_slug_kwd": slug, + } + ) + + before = len(chunks) + chunks = _narrow_by_keywords(chunks, keywords) + _LOG.info("[Wiki lookup] Found %d wiki page(s), kept %d after keyword filtering.", before, len(chunks)) + + doc_aggs: list[dict] = [] + seen: set = set() + for c in chunks: + did = c.get("doc_id") + if did and did not in seen: + seen.add(did) + doc_aggs.append({"doc_id": did, "doc_name": c.get("docnm_kwd") or ""}) + + return {"answer": "", "chunks": chunks, "doc_aggs": doc_aggs} diff --git a/rag/advanced_rag/harness/tools/gating.py b/rag/advanced_rag/harness/tools/gating.py index 3605c1a6ab..7b1f2485c9 100644 --- a/rag/advanced_rag/harness/tools/gating.py +++ b/rag/advanced_rag/harness/tools/gating.py @@ -25,12 +25,13 @@ SEARCH_PHASES = { "tools_priority": [ "hybrid_search", "bm25_search", + "web_search", "graph_explore", "inspector_open_context", "inspector_request_adjacent", ], "max_returned": 4, - "tool_hint": "Prefer retrieval tools to gather detailed information within the located region.", + "tool_hint": "Prefer retrieval tools to gather detailed information within the located region; use web_search when the knowledge base lacks the answer or the question needs current/external facts.", }, "verify": { "goal": "Verify consistency across multiple sources.", @@ -38,11 +39,11 @@ SEARCH_PHASES = { "inspector_open_context", "inspector_compare", "inspector_grep_within", - "hybrid_search", "web_search", + "hybrid_search", ], "max_returned": 4, - "tool_hint": "Prefer inspector tools to compare existing evidence before searching for new content.", + "tool_hint": "Prefer inspector tools to compare existing evidence; use web_search to corroborate against external sources.", }, "cross_domain": { "goal": "Explore cross-domain relationships for discovered entities.", @@ -89,16 +90,20 @@ def get_gated_tools( compilation_map: dict[str, set[str]], context: OrchestratorContext, has_routed_scope: bool = False, + web_enabled: bool = True, ) -> list[dict]: """Filter, sort, and gate tools by phase priority and context.""" phase_config = SEARCH_PHASES.get(phase) if not phase_config: - return _default_defs(available_tools) + return _default_defs(available_tools, web_enabled) sorted_tools = [] for tool_name in phase_config["tools_priority"]: if tool_name not in available_tools: continue + if tool_name == "web_search" and not web_enabled: + # No web provider configured — don't bind a tool that no-ops. + continue if not compilation_available(tool_name, compilation_map): continue if not tool_fits_context(tool_name, context, has_routed_scope): @@ -113,8 +118,8 @@ def get_gated_tools( return defs -def _default_defs(tool_names: list[str]) -> list[dict]: - return [TOOL_REGISTRY[n]["function_schema"] for n in tool_names if n in TOOL_REGISTRY] +def _default_defs(tool_names: list[str], web_enabled: bool = True) -> list[dict]: + return [TOOL_REGISTRY[n]["function_schema"] for n in tool_names if n in TOOL_REGISTRY and (web_enabled or n != "web_search")] def determine_current_phase(context: OrchestratorContext) -> str: diff --git a/rag/advanced_rag/harness/tools/search.py b/rag/advanced_rag/harness/tools/search.py index 4f26e974ab..492a6d6583 100644 --- a/rag/advanced_rag/harness/tools/search.py +++ b/rag/advanced_rag/harness/tools/search.py @@ -217,6 +217,7 @@ async def hybrid_search(tools, query: str, kb_ids: list[str] | None = None, top_ aggs=True, highlight=False, doc_ids=doc_scope, + must_not={"exists": "compile_kwd"}, # plain retrieval = document chunks only; compiled products have their own tools ) kbinfos = _normalize(kbinfos, tools.tenant_ids) if keywords: @@ -253,6 +254,7 @@ async def vector_search(tools, query: str, kb_ids: list[str] | None = None, top_ aggs=False, highlight=False, doc_ids=doc_scope, + must_not={"exists": "compile_kwd"}, ) kbinfos = _normalize(kbinfos, tools.tenant_ids) if keywords: @@ -280,6 +282,7 @@ async def bm25_search(tools, query: str, kb_ids: list[str] | None = None, top_n: aggs=False, highlight=False, doc_ids=doc_scope, + must_not={"exists": "compile_kwd"}, ) kbinfos = _normalize(kbinfos, tools.tenant_ids) if keywords: diff --git a/rag/advanced_rag/knowlege_compile/wiki.py b/rag/advanced_rag/knowlege_compile/wiki.py index d1a37c4c79..fc2437d3fb 100644 --- a/rag/advanced_rag/knowlege_compile/wiki.py +++ b/rag/advanced_rag/knowlege_compile/wiki.py @@ -3375,16 +3375,25 @@ async def _wiki_persist_draft( tenant_id: str, kb_id: str, plan_input_hash: str = "", + embd_mdl=None, ) -> None: - """Upsert one non-searchable wiki_page_draft row (resume cache). + """Upsert one wiki_page_draft row (resume cache + searchable page). ``plan_input_hash`` is the PLAN's ``input_hash_kwd`` at the time this draft was produced. The next REFINE re-entry compares it against the current PLAN hash to decide whether the cached draft is still valid; a mismatch forces a rewrite for that slug. + + When ``embd_mdl`` is provided the row is made searchable: the title/body are + tokenized (``title_tks`` / ``content_ltks`` / ``content_sm_ltks``) and a + ``q__vec`` page embedding is attached, with ``available_int=1`` so the + agent's ``wiki_query`` tool can retrieve it. Without an embedder the row stays + a non-searchable resume cache. ``content_with_weight`` is left as the page + JSON either way, so ``_wiki_load_refine_resume`` still restores the draft. """ from common import settings from rag.nlp import search as _rag_search + from rag.nlp import rag_tokenizer slug = page.get("slug") or "" if not slug: @@ -3401,8 +3410,37 @@ async def _wiki_persist_draft( "source_doc_ids": draft_doc_ids, "input_hash_kwd": plan_input_hash, "content_with_weight": content_with_weight, - "available_int": 0, # non-searchable + "available_int": 0, # non-searchable unless made searchable below } + + # Make the draft searchable when an embedder is available. content_with_weight + # is deliberately left untouched (the page JSON) — the tokenized fields drive + # BM25 and q__vec drives dense retrieval. + if embd_mdl is not None: + title = str(page.get("title") or slug) + body = str(page.get("content_md_rendered") or page.get("content_md") or page.get("content_md_raw") or "") + summary = str(page.get("summary") or "") + content_ltks = rag_tokenizer.tokenize(body) + row.update( + { + "docnm_kwd": title, + "title_kwd": title, + "title_tks": rag_tokenizer.tokenize(title), + "content_ltks": content_ltks, + "content_sm_ltks": rag_tokenizer.fine_grained_tokenize(content_ltks), + } + ) + try: + emb_text = (summary or f"{title}\n{body}").strip()[:2048] or title + vectors, _ = await thread_pool_exec(embd_mdl.encode, [emb_text]) + vec = vectors[0] + vec_list = vec.tolist() if hasattr(vec, "tolist") else list(vec) + if vec_list: + row[f"q_{len(vec_list)}_vec"] = vec_list + row["available_int"] = 1 + except Exception: + logging.exception("wiki_refine: draft embedding failed slug=%s; row stays non-searchable", slug) + try: try: await thread_pool_exec( @@ -3750,6 +3788,7 @@ async def wiki_refine_from_plan( tenant_id, kb_id, plan_input_hash=plan_input_hash, + embd_mdl=embd_mdl, ) except Exception: logging.exception("wiki_refine: persist_draft failed for slug=%s", slug) @@ -3820,6 +3859,7 @@ async def wiki_refine_from_plan( tenant_id, kb_id, plan_input_hash=plan_input_hash, + embd_mdl=embd_mdl, ) except Exception: logging.exception("wiki_refine: persist cleaned draft failed for slug=%s", page.get("slug")) diff --git a/rag/llm/chat_model.py b/rag/llm/chat_model.py index c19f89a3ec..9a4d972782 100644 --- a/rag/llm/chat_model.py +++ b/rag/llm/chat_model.py @@ -1583,6 +1583,53 @@ class GreenPTChat(Base): super().__init__(key, model_name, base_url or "https://api.greenpt.ai/v1", **kwargs) +# MiniMax models sometimes emit their bracket-delimited control/boundary tokens +# into `content` instead of as structured control — e.g. "]<]minimax[>[" — most +# often on tool-calling turns. The token is streamed split across many deltas, +# so it can't be removed per-delta; it must be filtered over a window that spans +# chunk boundaries. This pattern only matches the vendor name when it is wrapped +# in bracket noise on BOTH sides, so ordinary prose that mentions "MiniMax" is +# left untouched. Extend the alternation as further control tokens are observed. +_MINIMAX_CONTROL_TOKEN_RE = re.compile(r"[\[\]<>]+\s*minimax\s*[\[\]<>]+", re.IGNORECASE) + + +class _StreamSanitizer: + """Strip a regex from a token stream even when matches span chunk boundaries. + + A control token is bracket+letter characters, and it can arrive split across + many deltas, so we hold back the trailing run of token-ish characters (which + might still be forming a match) and only ``sub`` + emit the part before it. + Applying ``sub`` to a partial trailing run would fire prematurely and leak the + unmatched remainder — hence the hold. ``flush()`` sanitizes and returns the + remainder at end of stream. ``keep`` caps how long a run is buffered so a very + long separator-less word can't stall the stream forever. + """ + + _TOKENISH = re.compile(r"[\[\]<>A-Za-z]*$") + + def __init__(self, pattern: re.Pattern, keep: int = 64) -> None: + self._pat = pattern + self._keep = keep + self._buf = "" + + def feed(self, text: str) -> str: + if not text: + return "" + self._buf += text + match = self._TOKENISH.search(self._buf) + hold_start = match.start() if match else len(self._buf) + if len(self._buf) - hold_start > self._keep: + hold_start = len(self._buf) - self._keep + emit = self._pat.sub("", self._buf[:hold_start]) + self._buf = self._buf[hold_start:] + return emit + + def flush(self) -> str: + out = self._pat.sub("", self._buf) + self._buf = "" + return out + + class LiteLLMBase(ABC): _FACTORY_NAME = [ "Tongyi-Qianwen", @@ -1722,6 +1769,18 @@ class LiteLLMBase(ABC): def _need_reasoning_content_back(self) -> bool: return self.provider == SupportedLiteLLMProvider.DeepSeek + def _content_stream_sanitizer(self) -> "_StreamSanitizer | None": + """A per-stream filter for providers whose control tokens leak into content.""" + if self.provider == SupportedLiteLLMProvider.MiniMax: + return _StreamSanitizer(_MINIMAX_CONTROL_TOKEN_RE) + return None + + def _sanitize_answer(self, text: str) -> str: + """Strip provider control-token noise from a fully-assembled answer.""" + if text and self.provider == SupportedLiteLLMProvider.MiniMax: + return _MINIMAX_CONTROL_TOKEN_RE.sub("", text) + return text + async def async_chat(self, system, history, gen_conf, **kwargs): hist = list(history) if history else [] if system: @@ -2021,7 +2080,7 @@ class LiteLLMBase(ABC): ans += message.content or "" if response.choices[0].finish_reason == "length": ans = self._length_stop(ans) - return ans, tk_count + return self._sanitize_answer(ans), tk_count async def _exec_tool(tc): name = tc.function.name @@ -2060,7 +2119,7 @@ class LiteLLMBase(ABC): agg_usage["total_tokens"] += int(_fb.get("total_tokens", 0) or token_count) tk_count = agg_usage["total_tokens"] self.last_usage = dict(agg_usage) - return ans, tk_count + return self._sanitize_answer(ans), tk_count except Exception as e: e = await self._exceptions_async(e, attempt) @@ -2115,6 +2174,9 @@ class LiteLLMBase(ABC): answer = "" round_usage = None round_estimate = 0 + # Per-round filter for providers (MiniMax) whose control tokens + # leak into content split across deltas; None for others. + _sanitizer = self._content_stream_sanitizer() async for resp in response: # Usage-only final chunk may carry no choices — read it first. @@ -2154,7 +2216,12 @@ class LiteLLMBase(ABC): else: reasoning_start = False answer += delta.content - yield delta.content + if _sanitizer is not None: + emitted = _sanitizer.feed(delta.content) + if emitted: + yield emitted + else: + yield delta.content if not _u["total_tokens"]: round_estimate += num_tokens_from_string(delta.content) @@ -2163,6 +2230,12 @@ class LiteLLMBase(ABC): if finish_reason == "length": yield self._length_stop("") + # Flush any held-back (sanitized) answer content for this round. + if _sanitizer is not None: + tail = _sanitizer.flush() + if tail: + yield tail + # Commit this round's tokens to the running aggregate. _commit_round(round_usage, round_estimate) diff --git a/rag/nlp/search.py b/rag/nlp/search.py index 32940fdff5..50d2462903 100644 --- a/rag/nlp/search.py +++ b/rag/nlp/search.py @@ -563,6 +563,7 @@ class Dealer: highlight=False, rank_feature: dict | None = {PAGERANK_FLD: 10}, trace_id=None, + must_not: dict | None = None, ): ranks = {"total": 0, "chunks": [], "doc_aggs": {}} if not question: @@ -587,6 +588,8 @@ class Dealer: "similarity": similarity_threshold, "available_int": 1, } + if isinstance(must_not, dict) and must_not: + req["must_not"] = must_not logging.debug(f"[Search] global_offset={global_offset}, rerank_limit={RERANK_LIMIT}, page_size={page_size}, page={page}") if isinstance(tenant_ids, str):