392 lines
15 KiB
Python
392 lines
15 KiB
Python
"""本地知识库 — 加载 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
|
||
|
||
def search(self, query: str, top_k: int = 3, min_score: float = 0.05) -> 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
|
||
|
||
# 按文档聚合：取各文档下的最高 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)
|
||
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 []
|
||
|
||
# 文档级排序：主分(最高 section 分) + 0.02 * section 数量（打破单个高分 lucky hit）
|
||
sorted_docs = sorted(
|
||
doc_best.items(),
|
||
key=lambda x: x[1] + 0.02 * 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 = {"版本号", "版本", "区分", "定义"}
|
||
|
||
# 定义类查询：检查每个文档是否有定义类 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
|
||
if has_action and len(doc_sections) > 1:
|
||
# 动作查询：按优先级匹配 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)
|