@hupan56/wlkj 3.0.0 → 3.1.2

package/templates/qoder/scripts/{common → domain/kg/search}/search_engine.py

@@ -1,205 +1,207 @@
-# -*- coding: utf-8 -*-
-"""
-search_engine.py - 搜索倒排索引 + mtime 缓存 (性能优化 A2)
-
-把 search_keywords 的 O(n) 全扫描 3267 keys 降到 O(命中数)。
-
-设计:
-    - exact_index: {word_lower: [keyword_keys]} 精确匹配, O(1) 查找
-    - prefix/suffix 索引用 trigram (3-gram) 加速子串匹配
-    - 缓存序列化到 data/index/.inverted-cache.json, 含源 mtime
-    - 只在 code-keyword.json 变化时重建
-
-核心 API:
-    get_inverted_index(repo_root) -> dict   # 带缓存的倒排索引
-    search_keywords_fast(words, inverted, ki) -> {kw: files}  # 快速搜索
-"""
-
-import json
-import os
-import sys
-from typing import Dict, List, Optional, Set, Tuple
-
-_THIS_DIR = os.path.dirname(os.path.abspath(__file__))
-if _THIS_DIR not in sys.path:
-    sys.path.insert(0, _THIS_DIR)
-
-from atomicio import safe_read_json, atomic_write_json
-
-__all__ = [
-    "get_inverted_index",
-    "search_keywords_fast",
-    "build_inverted_index",
-    "MIN_FUZZY_LEN",
-]
-
-MIN_FUZZY_LEN = 4  # 与 search_index.py 保持一致
-
-
-def _trigrams(s: str) -> Set[str]:
-    """生成字符串的所有 3-gram (小写)。短于 3 字符则返回 {字符串本身}。"""
-    s = s.lower()
-    if len(s) < 3:
-        return {s}
-    return {s[i:i+3] for i in range(len(s) - 2)}
-
-
-def build_inverted_index(keyword_index: Dict[str, list]) -> dict:
-    """从 keyword_index 构建倒排索引。
-
-    Returns:
-        {
-            "exact": {word_lower: [keyword_keys]},      # 精确匹配
-            "trigram": {trigram: set(keyword_keys)},    # 3-gram → 含它的 keys
-            "keywords": [所有 keyword_keys 排序]         # 用于回退全扫描
-        }
-    """
-    exact = {}
-    trigram = {}
-    all_keys = sorted(keyword_index.keys())
-
-    for kw in all_keys:
-        kl = kw.lower()
-        # 精确索引
-        exact.setdefault(kl, []).append(kw)
-        # trigram 索引 (用于子串匹配的候选集缩小)
-        for tg in _trigrams(kl):
-            trigram.setdefault(tg, set()).add(kw)
-
-    # set 转 list 便于 JSON 序列化
-    trigram_serializable = {tg: sorted(keys) for tg, keys in trigram.items()}
-
-    return {
-        "exact": exact,
-        "trigram": trigram_serializable,
-        "keywords": all_keys,
-    }
-
-
-def get_inverted_index(repo_root: str) -> Optional[dict]:
-    """获取倒排索引 (带 mtime 缓存)。
-
-    缓存逻辑:
-    - 读 .inverted-cache.json 的 source_mtime
-    - 与 code-keyword.json 当前 mtime 比对
-    - 不一致或缓存缺失则重建 + 持久化
-
-    Returns:
-        倒排索引 dict, 或 None (keyword-index 不存在时)。
-    """
-    index_dir = os.path.join(repo_root, "data", "index")
-    ki_path = os.path.join(index_dir, "code-keyword.json")
-    cache_path = os.path.join(index_dir, ".inverted-cache.json")
-
-    if not os.path.isfile(ki_path):
-        return None
-
-    try:
-        ki_mtime = os.path.getmtime(ki_path)
-    except OSError:
-        return None
-
-    # 尝试读缓存
-    cached = safe_read_json(cache_path, default=None)
-    if cached and cached.get("source_mtime") == ki_mtime:
-        # 缓存命中 (trigram 的 set 反序列化为 list, 搜索时兼容)
-        return cached
-
-    # 缓存失效或缺失, 重建
-    ki = safe_read_json(ki_path, default=None)
-    if not ki:
-        return None
-
-    inverted = build_inverted_index(ki)
-    inverted["source_mtime"] = ki_mtime
-    inverted["built_at"] = os.path.getmtime(ki_path)
-
-    # 持久化 (best-effort, 失败不阻塞搜索)
-    try:
-        atomic_write_json(cache_path, inverted)
-    except Exception:
-        pass
-
-    return inverted
-
-
-def _candidate_keys(word: str, inverted: dict) -> List[str]:
-    """用倒排索引找到候选 keyword keys (缩小扫描集)。
-
-    策略:
-    - 用 trigram 交集找候选 (子串匹配的双方至少共享一个 trigram)
-    - 精确命中必然包含在 trigram 候选里 (相同字符串共享所有 trigram)
-    - 候选为空则回退全部 keys (兜底, 罕见)
-
-    关键: 不能在精确命中后短路 —— query "insurance" 还要匹配
-    "insurancepolicy" 等 (query 是 kw 的子串), 所以必须返回 trigram 候选集。
-    """
-    wl = word.lower()
-    tgs = _trigrams(wl)
-    trigram_index = inverted.get("trigram", {})
-    candidate_sets = []
-    for tg in tgs:
-        keys = trigram_index.get(tg)
-        if keys:
-            candidate_sets.append(set(keys))
-
-    if not candidate_sets:
-        # 词太短无 trigram (如 2 字符), 回退精确 + 全扫描兜底
-        exact = inverted.get("exact", {}).get(wl)
-        if exact:
-            return list(exact)
-        return inverted.get("keywords", [])
-
-    # 取并集: query 含 kw 或 kw 含 query, 至少一个 trigram 重叠
-    candidates = set()
-    for cs in candidate_sets:
-        candidates |= cs
-
-    return sorted(candidates)
-
-
-def search_keywords_fast(
-    words: List[str],
-    inverted: Optional[dict],
-    keyword_index: Dict[str, list],
-) -> Dict[str, list]:
-    """快速关键词搜索 (用倒排索引缩小候选集)。
-
-    Args:
-        words: 查询词列表 (已小写化、已 CN 扩展)。
-        inverted: 倒排索引 (None 则回退到全扫描)。
-        keyword_index: 原始 keyword-index (用于取 files)。
-
-    Returns:
-        {keyword_key: files} 匹配结果。
-    """
-    matches = {}
-
-    if inverted is None:
-        # 回退: 全扫描 (旧逻辑)
-        for word in words:
-            for kw, files in keyword_index.items():
-                if _match(word, kw):
-                    matches[kw] = files
-        return matches
-
-    for word in words:
-        candidates = _candidate_keys(word, inverted)
-        for kw in candidates:
-            files = keyword_index.get(kw)
-            if files and _match(word, kw):
-                matches[kw] = files
-
-    return matches
-
-
-def _match(query_word: str, kw: str) -> bool:
-    """与 search_index._match_keyword 一致的匹配逻辑。"""
-    w, k = query_word.lower(), kw.lower()
-    if w == k:
-        return True
-    if len(k) >= MIN_FUZZY_LEN and k in w:
-        return True
-    if len(w) >= MIN_FUZZY_LEN and w in k:
-        return True
-    return False
+# -*- coding: utf-8 -*-
+"""
+search_engine.py - 搜索倒排索引 + mtime 缓存 (性能优化 A2)
+
+把 search_keywords 的 O(n) 全扫描 3267 keys 降到 O(命中数)。
+
+设计:
+    - exact_index: {word_lower: [keyword_keys]} 精确匹配, O(1) 查找
+    - prefix/suffix 索引用 trigram (3-gram) 加速子串匹配
+    - 缓存序列化到 data/index/.inverted-cache.json, 含源 mtime
+    - 只在 code-keyword.json 变化时重建
+
+核心 API:
+    get_inverted_index(repo_root) -> dict   # 带缓存的倒排索引
+    search_keywords_fast(words, inverted, ki) -> {kw: files}  # 快速搜索
+"""
+
+import json
+import os
+import sys
+from typing import Dict, List, Optional, Set, Tuple
+
+# v3.0 路径自举: 迁入 kg/ 后, atomicio 在 common/, 需把 common/ 加进 sys.path
+_THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+_cp = os.path.join(os.path.dirname(_THIS_DIR), 'foundation')
+if _cp not in sys.path:
+    sys.path.insert(0, _cp)
+
+from atomicio import safe_read_json, atomic_write_json
+
+__all__ = [
+    "get_inverted_index",
+    "search_keywords_fast",
+    "build_inverted_index",
+    "MIN_FUZZY_LEN",
+]
+
+MIN_FUZZY_LEN = 4  # 与 search_index.py 保持一致
+
+
+def _trigrams(s: str) -> Set[str]:
+    """生成字符串的所有 3-gram (小写)。短于 3 字符则返回 {字符串本身}。"""
+    s = s.lower()
+    if len(s) < 3:
+        return {s}
+    return {s[i:i+3] for i in range(len(s) - 2)}
+
+
+def build_inverted_index(keyword_index: Dict[str, list]) -> dict:
+    """从 keyword_index 构建倒排索引。
+
+    Returns:
+        {
+            "exact": {word_lower: [keyword_keys]},      # 精确匹配
+            "trigram": {trigram: set(keyword_keys)},    # 3-gram → 含它的 keys
+            "keywords": [所有 keyword_keys 排序]         # 用于回退全扫描
+        }
+    """
+    exact = {}
+    trigram = {}
+    all_keys = sorted(keyword_index.keys())
+
+    for kw in all_keys:
+        kl = kw.lower()
+        # 精确索引
+        exact.setdefault(kl, []).append(kw)
+        # trigram 索引 (用于子串匹配的候选集缩小)
+        for tg in _trigrams(kl):
+            trigram.setdefault(tg, set()).add(kw)
+
+    # set 转 list 便于 JSON 序列化
+    trigram_serializable = {tg: sorted(keys) for tg, keys in trigram.items()}
+
+    return {
+        "exact": exact,
+        "trigram": trigram_serializable,
+        "keywords": all_keys,
+    }
+
+
+def get_inverted_index(repo_root: str) -> Optional[dict]:
+    """获取倒排索引 (带 mtime 缓存)。
+
+    缓存逻辑:
+    - 读 .inverted-cache.json 的 source_mtime
+    - 与 code-keyword.json 当前 mtime 比对
+    - 不一致或缓存缺失则重建 + 持久化
+
+    Returns:
+        倒排索引 dict, 或 None (keyword-index 不存在时)。
+    """
+    index_dir = os.path.join(repo_root, "data", "index")
+    ki_path = os.path.join(index_dir, "code-keyword.json")
+    cache_path = os.path.join(index_dir, ".inverted-cache.json")
+
+    if not os.path.isfile(ki_path):
+        return None
+
+    try:
+        ki_mtime = os.path.getmtime(ki_path)
+    except OSError:
+        return None
+
+    # 尝试读缓存
+    cached = safe_read_json(cache_path, default=None)
+    if cached and cached.get("source_mtime") == ki_mtime:
+        # 缓存命中 (trigram 的 set 反序列化为 list, 搜索时兼容)
+        return cached
+
+    # 缓存失效或缺失, 重建
+    ki = safe_read_json(ki_path, default=None)
+    if not ki:
+        return None
+
+    inverted = build_inverted_index(ki)
+    inverted["source_mtime"] = ki_mtime
+    inverted["built_at"] = os.path.getmtime(ki_path)
+
+    # 持久化 (best-effort, 失败不阻塞搜索)
+    try:
+        atomic_write_json(cache_path, inverted)
+    except Exception:
+        pass
+
+    return inverted
+
+
+def _candidate_keys(word: str, inverted: dict) -> List[str]:
+    """用倒排索引找到候选 keyword keys (缩小扫描集)。
+
+    策略:
+    - 用 trigram 交集找候选 (子串匹配的双方至少共享一个 trigram)
+    - 精确命中必然包含在 trigram 候选里 (相同字符串共享所有 trigram)
+    - 候选为空则回退全部 keys (兜底, 罕见)
+
+    关键: 不能在精确命中后短路 —— query "insurance" 还要匹配
+    "insurancepolicy" 等 (query 是 kw 的子串), 所以必须返回 trigram 候选集。
+    """
+    wl = word.lower()
+    tgs = _trigrams(wl)
+    trigram_index = inverted.get("trigram", {})
+    candidate_sets = []
+    for tg in tgs:
+        keys = trigram_index.get(tg)
+        if keys:
+            candidate_sets.append(set(keys))
+
+    if not candidate_sets:
+        # 词太短无 trigram (如 2 字符), 回退精确 + 全扫描兜底
+        exact = inverted.get("exact", {}).get(wl)
+        if exact:
+            return list(exact)
+        return inverted.get("keywords", [])
+
+    # 取并集: query 含 kw 或 kw 含 query, 至少一个 trigram 重叠
+    candidates = set()
+    for cs in candidate_sets:
+        candidates |= cs
+
+    return sorted(candidates)
+
+
+def search_keywords_fast(
+    words: List[str],
+    inverted: Optional[dict],
+    keyword_index: Dict[str, list],
+) -> Dict[str, list]:
+    """快速关键词搜索 (用倒排索引缩小候选集)。
+
+    Args:
+        words: 查询词列表 (已小写化、已 CN 扩展)。
+        inverted: 倒排索引 (None 则回退到全扫描)。
+        keyword_index: 原始 keyword-index (用于取 files)。
+
+    Returns:
+        {keyword_key: files} 匹配结果。
+    """
+    matches = {}
+
+    if inverted is None:
+        # 回退: 全扫描 (旧逻辑)
+        for word in words:
+            for kw, files in keyword_index.items():
+                if _match(word, kw):
+                    matches[kw] = files
+        return matches
+
+    for word in words:
+        candidates = _candidate_keys(word, inverted)
+        for kw in candidates:
+            files = keyword_index.get(kw)
+            if files and _match(word, kw):
+                matches[kw] = files
+
+    return matches
+
+
+def _match(query_word: str, kw: str) -> bool:
+    """与 search_index._match_keyword 一致的匹配逻辑。"""
+    w, k = query_word.lower(), kw.lower()
+    if w == k:
+        return True
+    if len(k) >= MIN_FUZZY_LEN and k in w:
+        return True
+    if len(w) >= MIN_FUZZY_LEN and w in k:
+        return True
+    return False