"""本地知识库 — 加载 Markdown/TXT 文件，jieba 分词 + TF-IDF 搜索 核心改进：知识库按标题（## / ###）切分为独立的 section，每个 section 作为 一个独立的"文档"进行索引。搜索时直接找到最相关的 section，返回结果即为 最相关的片段，不再需要从整篇文档中二次提取。 """ import os import re import logging from pathlib import Path from typing import List, Tuple, Optional import jieba from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.metrics.pairwise import cosine_similarity import numpy as np class KnowledgeBase: """知识库引擎""" def __init__(self, root_path: str = ""): self.root_path: str = root_path self.docs: List[dict] = [] # [{path, title, content, heading, section_id}] self._vectorizer = TfidfVectorizer(tokenizer=self._tokenize, lowercase=False) self._tfidf_matrix = None self._ready = False self._heading_token_sets: List[set] = [] self._body_token_sets: List[set] = [] # 常见同义词映射（查询侧 → 文档侧） SYNONYMS = { "登陆": ["登录"], "版号": ["版本号"], "权力": ["权利"], } @staticmethod def _tokenize(text: str) -> List[str]: """jieba 分词""" return list(jieba.cut(text)) @classmethod def _expand_query(cls, query: str) -> List[str]: """扩展查询：为每个词生成同义词变体，提高对错别字/同义词的容忍度。 返回多个候选查询字符串。 """ # 基础查询 queries = [query] # 为每个同义词替换生成变体 for src, dests in cls.SYNONYMS.items(): if src in query: for dest in dests: queries.append(query.replace(src, dest)) return queries @staticmethod def _parse_qa_segments(content: str) -> list: """解析 Q&A 格式文件为独立问答段。 格式：问题[N]：xxx\n答：xxx 返回 [{"question": str, "answer": str, "raw": str}] """ segments = [] cur_q = [] cur_a = [] state = None # None, 'q', 'a' def flush(): if not cur_q and not cur_a: return q = '\n'.join(cur_q).strip() a = '\n'.join(cur_a).strip() if q or a: segments.append({ "question": q, "answer": a, "raw": f"问题：{q}\n答：{a}", }) cur_q.clear() cur_a.clear() for line in content.split('\n'): qm = re.match(r'^问题\d*\s*[：:]\s*(.*)', line) am = re.match(r'^答\s*[：:]\s*(.*)', line) if qm: flush() state = 'q' cur_q.append(qm.group(1)) elif am: state = 'a' cur_a.append(am.group(1)) elif state == 'q': cur_q.append(line) elif state == 'a': cur_a.append(line) # else: 文件头空白/注释，忽略 flush() return segments @staticmethod def _split_into_sections(content: str, title: str) -> List[dict]: """把 Markdown 文档按 ## / ### 标题切分为多个 section。 每个 section 包含： - heading: 标题行（如 "## 可用参数"） - content: 标题 + 正文（用于展示） - body: 仅正文（用于 TF-IDF，避免标题词重复加权） 返回 section 列表。 """ sections = [] # 按行首的 ## 标题切分（前面可能有空白） parts = re.split(r'(?=^#{1,6}\s)', content, flags=re.MULTILINE) for i, part in enumerate(parts): part_stripped = part.strip() if not part_stripped or len(part_stripped) < 5: continue # 提取标题行 lines = part_stripped.split('\n') heading = lines[0].strip() body = '\n'.join(lines[1:]).strip() if len(lines) > 1 else "" sections.append({ "heading": heading, "content": part_stripped, "body": body, "section_id": i, }) # 如果没有切分出多个 section，整个文档作为一段 if len(sections) <= 1: sections = [{ "heading": f"# {title}", "content": content.strip(), "body": content.strip(), "section_id": 0, }] return sections def set_path(self, path: str): self.root_path = path self.reload() def reload(self): """重新加载知识库文件""" self.docs = [] if not self.root_path or not os.path.isdir(self.root_path): self._ready = False return root = Path(self.root_path) extensions = (".md", ".txt", ".markdown") for fp in sorted(root.rglob("*")): if fp.suffix.lower() in extensions and fp.is_file(): try: content = fp.read_text("utf-8", errors="ignore") title = fp.stem # 解析 Q&A 段 segments = self._parse_qa_segments(content) if segments: for seg in segments: self.docs.append({ "path": str(fp), "title": title, "content": seg["raw"], "answer": seg["answer"], "question": seg["question"], "heading": f"# {title}", "body": seg["raw"], }) else: # 非 Q&A 格式 → 按标题切分为多个 section sections = self._split_into_sections(content, title) for sec in sections: self.docs.append({ "path": str(fp), "title": title, "content": sec["content"], "heading": sec["heading"], "body": sec["body"], }) except Exception: continue if not self.docs: self._ready = False return # 构建 TF-IDF（用 body 而非 content，避免标题词重复加权） corpus = [d.get("body", d["content"]) for d in self.docs] try: self._tfidf_matrix = self._vectorizer.fit_transform(corpus) self._ready = True except Exception: self._ready = False return # 预计算分词集合，避免每次搜索重复 jieba self._heading_token_sets = [] self._body_token_sets = [] for doc in self.docs: heading = doc.get("heading", "") body = doc.get("body", doc["content"]) self._heading_token_sets.append(set(t for t in jieba.cut(heading) if len(t) >= 2)) self._body_token_sets.append(set(t for t in jieba.cut(body) if len(t) >= 2)) @property def is_ready(self) -> bool: return self._ready and len(self.docs) > 0 @staticmethod def _extract_verb_object(query: str, action_verbs: set) -> List[Tuple[str, str]]: """从查询中提取动词+对象对。 例："怎么打印流水" → [("打印", "流水")] "圈地工具用什么" → [("圈地", "工具")] """ import re pairs = [] # 模式1: 动词 + 对象（2+字动词 + 2+字名词） for verb in action_verbs: if verb in query: # 找动词后面的内容作为对象 idx = query.index(verb) + len(verb) rest = query[idx:].strip() # 取动词后 2~6 字作为对象 if len(rest) >= 2: obj = rest[:6].strip() # 去掉常见无意义尾词 for stop in ["的", "地", "得", "了", "吗", "呢", "吧", "啊"]: obj = obj.rstrip(stop) if len(obj) >= 2: pairs.append((verb, obj)) # 模式2: 2字动词 + 2字对象（从分词结果匹配） if not pairs: tokens = [t for t in jieba.cut(query) if len(t) >= 2] for i, t in enumerate(tokens): if t in action_verbs and i + 1 < len(tokens): pairs.append((t, tokens[i + 1])) return pairs def search(self, query: str, top_k: int = 3, min_score: float = 0.1) -> List[Tuple[dict, float]]: """搜索知识库，返回 (doc, score) 列表。 策略： 1. TF-IDF + 关键词命中计算每个 section 的分数 2. 按文档聚合：找到最佳匹配的文档 3. 在最佳文档内，根据查询意图找到最合适的 section 4. 补充其他文档的高分 section """ if not self.is_ready or not query.strip(): return [] query_stripped = query.strip() # 扩展查询（同义词替换） expanded_queries = self._expand_query(query_stripped) # 短文本检查 tokens = [t for t in jieba.cut(query_stripped) if len(t) >= 2] if len(query_stripped) < 2 or len(tokens) == 0: return [] # 对每个扩展查询计算 TF-IDF，取最高分 tfidf_scores = np.zeros(len(self.docs)) for eq in expanded_queries: query_vec = self._vectorizer.transform([eq]) scores = cosine_similarity(query_vec, self._tfidf_matrix)[0] tfidf_scores = np.maximum(tfidf_scores, scores) # 关键词命中 expanded_tokens = set() for eq in expanded_queries: expanded_tokens.update(t for t in jieba.cut(eq) if len(t) >= 2) query_tokens_set = expanded_tokens idf_dict = {} if hasattr(self._vectorizer, 'idf_') and self._vectorizer.idf_ is not None: feature_names = self._vectorizer.get_feature_names_out() for t in query_tokens_set: if t in feature_names: idx = list(feature_names).index(t) idf_dict[t] = self._vectorizer.idf_[idx] else: idf_dict[t] = 1.0 else: for t in query_tokens_set: idf_dict[t] = 1.0 heading_token_sets = self._heading_token_sets keyword_scores = np.zeros(len(self.docs)) for i, doc in enumerate(self.docs): heading_hits = sum(idf_dict.get(t, 1.0) for t in query_tokens_set if t in heading_token_sets[i]) body_hits = sum(idf_dict.get(t, 1.0) for t in query_tokens_set if t in self._body_token_sets[i]) keyword_scores[i] = heading_hits * 3 + body_hits combined_scores = 0.7 * tfidf_scores + 0.3 * keyword_scores # 内容质量调整：惩罚标题-only/正文过短的 section，奖励正文含查询词的 section # 避免"# Sdf1系列登录插件使用说明"这种纯标题段因为标题词匹配就排第一 for i, doc in enumerate(self.docs): body = doc.get("body", "") body_len = len(body.strip()) heading = doc.get("heading", "") # 单层 # 标题且无实质正文 → 大幅降权 if heading.startswith("# ") and body_len < 50: combined_scores[i] *= 0.2 elif body_len < 30: combined_scores[i] *= 0.5 # 正文含查询词 → 奖励（说明 section 真的在讲用户问的内容） body_hits = sum(1 for t in query_tokens_set if t in self._body_token_sets[i]) if body_hits >= 2: combined_scores[i] *= 1.3 elif body_hits >= 1: combined_scores[i] *= 1.1 # 短查询惩罚：分词数 ≤2 的短查询（如图标残片/短词）得分大幅降权 # 防止 OCR 碎片中的高 IDF 词（如"下载"）获得超过真实用户消息的分数 n_tokens = len(tokens) if n_tokens <= 2: combined_scores *= 0.3 elif n_tokens <= 3: combined_scores *= 0.6 # 按文档聚合：取各文档下的最高 section 分作为文档分 # 同时统计匹配 section 数量，用于后续排序（section 多的文档更可能是主文档） doc_best = {} doc_count = {} for i, doc in enumerate(self.docs): t = doc["title"] score = combined_scores[i] if score >= min_score: doc_best[t] = max(doc_best.get(t, 0), score) # 只统计强匹配 section（≥2×min_score），弱匹配不计入 bonus if score >= min_score * 2: doc_count[t] = doc_count.get(t, 0) + 1 else: doc_best[t] = max(doc_best.get(t, 0), score) if not doc_best: return [] # 文档级排序：主分优先，doc_count 仅做 tiebreaker（section 数更多→多角度匹配） # 弱匹配（score < min_score*2）不计入 doc_count，防低分堆叠 sorted_docs = sorted( doc_best.items(), key=lambda x: (x[1], doc_count.get(x[0], 0)), reverse=True ) results = [] # 查询意图检测 query_lower = {t.lower() for t in query_tokens_set} action_words = {"怎么", "如何", "用法", "教程", "操作"} def_words = {"版本", "版号", "版本号", "定义", "区分"} has_action = bool(action_words & query_lower) has_def = bool(def_words & query_lower) def_headings = {"版本号", "版本", "区分", "定义"} # 动作查询：提取动词+对象，对不含两者的文档降权 # 防止"圈地工具"匹配到含"工具"但不含"圈地"的文档（如寻宝插件） action_verbs_all = {"打印", "查看", "设置", "创建", "使用", "查询", "导出", "执行", "安装", "改", "修改", "更改", "更换", "重置"} verb_obj_pairs = self._extract_verb_object(query_stripped, action_verbs_all) if verb_obj_pairs: adjusted = [] for doc_title, doc_score in sorted_docs: # 检查文档内容是否包含动词+对象 doc_text = "" for i, doc in enumerate(self.docs): if doc["title"] == doc_title: doc_text += doc.get("heading", "") + " " + doc.get("content", "") # 计算匹配度：动词+对象都包含 → 满分；只包含一个 → 降权；都不包含 → 大幅降权 match_count = 0 for verb, obj in verb_obj_pairs: v_in = verb in doc_text o_in = obj in doc_text if v_in and o_in: match_count += 2 elif v_in or o_in: match_count += 1 if match_count == 0: doc_score *= 0.2 # 完全不匹配，大幅降权 elif match_count == 1: doc_score *= 0.5 # 部分匹配，降权 adjusted.append((doc_title, doc_score)) adjusted.sort(key=lambda x: x[1], reverse=True) sorted_docs = adjusted # 定义类查询：检查每个文档是否有定义类 section，没有则降权 if has_def: adjusted = [] for doc_title, doc_score in sorted_docs: has_def_section = False for i, doc in enumerate(self.docs): if doc["title"] == doc_title and combined_scores[i] >= min_score / 2: if def_headings & self._heading_token_sets[i]: has_def_section = True break if not has_def_section: doc_score *= 0.3 # 没有定义类 section 的文档大幅降权 adjusted.append((doc_title, doc_score)) # 重新排序 adjusted.sort(key=lambda x: x[1], reverse=True) sorted_docs = adjusted # 从最佳文档开始，逐个文档取 section for doc_title, doc_score in sorted_docs: if len(results) >= top_k: break # 收集该文档的所有 section doc_sections = [] for i, doc in enumerate(self.docs): if doc["title"] == doc_title and combined_scores[i] >= min_score: doc_sections.append((i, doc, combined_scores[i])) doc_sections.sort(key=lambda x: x[2], reverse=True) if not doc_sections: continue # 同文档内，根据查询意图优先选合适的 section # 1. heading 包含查询关键词 → bonus；文档头部（仅含文档名）→ penalty query_tokens = set(t for t in jieba.cut(query_stripped) if len(t) >= 2) reranked = [] for idx, doc, score in doc_sections: heading = doc.get("heading", "") body = doc.get("body", doc.get("content", "")) h_tokens = set(t for t in jieba.cut(heading) if len(t) >= 2) # heading 包含查询关键词 → 适度 bonus（2.0 每命中词） heading_hits = len(query_tokens & h_tokens) bonus = heading_hits * 2.0 # heading 包含查询连续短语 → 额外 bonus for qlen in range(min(len(query_stripped), 8), 3, -1): for start in range(len(query_stripped) - qlen + 1): phrase = query_stripped[start:start+qlen] if len(phrase) >= 4 and phrase in heading: bonus += 6.0 break # body 包含查询连续短语 → 强 bonus（说明正文真的在讲用户问的内容） for qlen in range(min(len(query_stripped), 8), 3, -1): for start in range(len(query_stripped) - qlen + 1): phrase = query_stripped[start:start+qlen] if len(phrase) >= 4 and phrase in body: bonus += 10.0 break # 文档头部惩罚：heading 是单层 # + 文档名（如 "# php管理员后端"），且无 body if heading.startswith("# ") and (not body or len(body) < 50): bonus -= 15.0 reranked.append((idx, doc, score + bonus)) reranked.sort(key=lambda x: x[2], reverse=True) doc_sections = reranked if has_action and len(doc_sections) > 1: # 动作查询：提取动词+对象，精准匹配 section # 如"怎么打印流水" → verb="打印", obj="流水"，优先选 heading 同时含两者的 section action_verbs = {"打印", "查看", "设置", "创建", "使用", "查询", "导出", "执行", "改", "修改", "更改", "更换", "重置"} verb_obj_pairs = self._extract_verb_object(query_stripped, action_verbs) if verb_obj_pairs: # 有动词+对象：找 heading 或 body 同时包含两者的 section best_idx = -1 best_match = 0 for i in range(len(doc_sections)): idx = doc_sections[i][0] heading = self.docs[idx].get("heading", "") content = self.docs[idx].get("content", "") for verb, obj in verb_obj_pairs: v_in_h = verb in heading or verb in content o_in_h = obj in heading or obj in content match_score = int(v_in_h) + int(o_in_h) if match_score > best_match: best_match = match_score best_idx = i if best_idx > 0: doc_sections.insert(0, doc_sections.pop(best_idx)) else: # 无动词+对象：按优先级匹配 heading 词（旧逻辑） # "参数/指令/使用" > "方法" > "安装" priority_order = [ ({"参数", "指令", "使用"}), # 最高优先级 ({"方法"}), # 中等优先级 ({"安装"}), # 最低优先级 ] for priority_level in priority_order: for i in range(len(doc_sections)): idx = doc_sections[i][0] if priority_level & self._heading_token_sets[idx]: doc_sections.insert(0, doc_sections.pop(i)) break else: continue break # 找到匹配就停止 if has_def and len(doc_sections) > 1: # 定义查询：选 heading 中匹配最多定义类词的 section best_idx = -1 best_count = 0 for i in range(len(doc_sections)): idx = doc_sections[i][0] match_count = len(def_headings & self._heading_token_sets[idx]) if match_count > best_count: best_count = match_count best_idx = i if best_idx > 0: doc_sections.insert(0, doc_sections.pop(best_idx)) # 从每个文档取前 1-2 个 section sub_k = 2 if len(results) < top_k - 1 else top_k - len(results) for i, (idx, doc, score) in enumerate(doc_sections): if i >= sub_k: break results.append((doc, float(score))) return results[:top_k] if results: return results # Fallback: 子串匹配（查询包含在文档标题或内容中） q = query_stripped if len(q) >= 4: for doc in self.docs: if q in doc["content"] or q in doc["title"]: logger = logging.getLogger("msgbot.kb") logger.debug(f"子串匹配: \"{q}\" → {doc['title']}") return [(doc, min_score)] return results def get_doc_count(self) -> int: return len(self.docs)