95 lines
3.0 KiB
Python
95 lines
3.0 KiB
Python
"""本地知识库 — 加载 Markdown/TXT 文件，jieba 分词 + TF-IDF 搜索"""
|
|
|
|
import os
|
|
import re
|
|
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, segments}]
|
|
self._vectorizer = TfidfVectorizer(tokenizer=self._tokenize, lowercase=False)
|
|
self._tfidf_matrix = None
|
|
self._ready = False
|
|
|
|
@staticmethod
|
|
def _tokenize(text: str) -> List[str]:
|
|
"""jieba 分词"""
|
|
return list(jieba.cut(text))
|
|
|
|
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
|
|
self.docs.append({
|
|
"path": str(fp),
|
|
"title": title,
|
|
"content": content,
|
|
})
|
|
except Exception:
|
|
continue
|
|
|
|
if not self.docs:
|
|
self._ready = False
|
|
return
|
|
|
|
# 构建 TF-IDF
|
|
corpus = [d["content"] for d in self.docs]
|
|
try:
|
|
self._tfidf_matrix = self._vectorizer.fit_transform(corpus)
|
|
self._ready = True
|
|
except Exception:
|
|
self._ready = False
|
|
|
|
@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.30) -> List[Tuple[dict, float]]:
|
|
"""搜索知识库，返回 (doc, score) 列表。"""
|
|
if not self.is_ready or not query.strip():
|
|
return []
|
|
|
|
# 短文本（< 4 个字或 <= 2 个有效词）不做搜索，避免噪音
|
|
tokens = [t for t in jieba.cut(query.strip()) if len(t) >= 2]
|
|
if len(query.strip()) < 4 or len(tokens) <= 1:
|
|
return []
|
|
|
|
query_vec = self._vectorizer.transform([query])
|
|
scores = cosine_similarity(query_vec, self._tfidf_matrix)[0]
|
|
top_indices = np.argsort(scores)[::-1][:top_k]
|
|
|
|
results = []
|
|
for idx in top_indices:
|
|
score = float(scores[idx])
|
|
if score < min_score:
|
|
continue
|
|
results.append((self.docs[idx], score))
|
|
return results
|
|
|
|
def get_doc_count(self) -> int:
|
|
return len(self.docs)
|