"""知识库管理模块 - RAG 检索增强生成 架构：文档 → 精细分块(代码块独立/标题层级切分) → 中文embedding(BGE) + numpy余弦相似度 → 关键词+向量混合检索 → LLM生成回答 """ import hashlib import json import os import re from pathlib import Path from urllib.parse import quote # ============ 配置 ============ KNOWLEDGE_DIR = Path(__file__).parent INDEX_DIR = KNOWLEDGE_DIR / "index" CACHE_FILE = KNOWLEDGE_DIR / "qa_cache.json" CACHE_VERSION = "v3-detail-grounded" DOCS_SOURCE = Path("/vol3/1000/Code/ngfchl.github.io/docs") OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434") OLLAMA_MODEL = os.environ.get("OLLAMA_MODEL", "qwen3:8b") SITE_URL = "https://ptools.fun" # 切分参数 CHUNK_SIZE = 900 CHUNK_OVERLAP = 120 TOP_K = 10 # 中文 embedding 模型 EMBED_MODEL = "BAAI/bge-small-zh-v1.5" # 需要跳过的文件 SKIP_FILES = {"agreement.md", "index.md", "password.md", "termnal.md", "ssh.md"} SKIP_SECTION_KEYWORDS = ["油猴", "tampermonkey", "greasemonkey", "userscript"] # ============ Embedding ============ _model = None def get_model(): global _model if _model is None: from sentence_transformers import SentenceTransformer cache_dir = str(KNOWLEDGE_DIR / "models") os.environ["CUDA_VISIBLE_DEVICES"] = "" print(f"📦 加载中文 embedding 模型: {EMBED_MODEL}") _model = SentenceTransformer(EMBED_MODEL, cache_folder=cache_dir) print("✅ 模型加载完成") return _model def embed_texts(texts: list[str]): model = get_model() return model.encode(texts, normalize_embeddings=True, show_progress_bar=True, batch_size=32, convert_to_numpy=True) # ============ URL 生成 ============ def doc_path_to_url(rel_path: str) -> str: html_path = rel_path.replace(".md", ".html") parts = html_path.split("/") encoded = "/".join(quote(p, safe="") for p in parts) return f"{SITE_URL}/{encoded}" # ============ 文档清洗（保留结构） ============ def clean_markdown(text: str) -> str: """清洗但保留代码块标记和链接""" text = re.sub(r'<[^>]+>', '', text) # 图片替换为 alt 文本 text = re.sub(r'!\[([^\]]*)\]\([^\)]+\)', r'[图片:\1]', text) # 链接保留文本+URL text = re.sub(r'\[([^\]]+)\]\(([^\)]+)\)', r'\1 (\2)', text) # frontmatter text = re.sub(r'^---\n.*?---\n', '', text, flags=re.DOTALL) text = re.sub(r'\n{3,}', '\n\n', text) return text.strip() # ============ 代码块检测 ============ def find_code_blocks(text: str) -> list[tuple[int, int]]: """找到所有代码块的起止位置""" blocks = [] lines = text.split('\n') in_block = False start = -1 for i, line in enumerate(lines): stripped = line.strip() if re.match(r'^```\w*$', stripped): if not in_block: in_block = True start = i else: blocks.append((start, i)) in_block = False return blocks # ============ 精细分块 ============ def split_document(text: str, rel_path: str, doc_title: str, doc_url: str) -> list[dict]: """ 精细分块策略： 1. 代码块作为独立 chunk（不拆分） 2. 每个标题层级下的内容独立成 chunk 3. 大段落按 CHUNK_SIZE 拆分 4. 每个 chunk 带上标题路径上下文 """ chunks = [] lines = text.split('\n') # 先找代码块位置 code_blocks = find_code_blocks(text) code_block_lines = set() for start, end in code_blocks: for i in range(start, end + 1): code_block_lines.add(i) # 按标题切分 sections = [] # [(level, heading, line_start, line_end)] current_heading_stack = [] # [(level, heading)] current_start = 0 for i, line in enumerate(lines): if i in code_block_lines: continue m = re.match(r'^(#{1,6})\s+(.+)', line.strip()) if m: level = len(m.group(1)) heading = m.group(2).strip() # 保存之前的 section if i > current_start: sections.append({ "level": current_heading_stack[-1][0] if current_heading_stack else 0, "heading": current_heading_stack[-1][1] if current_heading_stack else "", "path": [h for _, h in current_heading_stack], "line_start": current_start, "line_end": i - 1, }) # 更新 heading stack while current_heading_stack and current_heading_stack[-1][0] >= level: current_heading_stack.pop() current_heading_stack.append((level, heading)) current_start = i # 最后一个 section sections.append({ "level": current_heading_stack[-1][0] if current_heading_stack else 0, "heading": current_heading_stack[-1][1] if current_heading_stack else "", "path": [h for _, h in current_heading_stack], "line_start": current_start, "line_end": len(lines) - 1, }) chunk_id = 0 for sec in sections: sec_lines = lines[sec["line_start"]:sec["line_end"] + 1] sec_text = '\n'.join(sec_lines).strip() if len(sec_text) < 15: continue # 检查是否包含油猴关键词 if any(kw in sec_text.lower() for kw in SKIP_SECTION_KEYWORDS): continue # 构建标题路径 title_path = " > ".join(sec["path"]) if sec["path"] else doc_title if title_path and title_path != doc_title: title_path = f"{doc_title} > {title_path}" else: title_path = doc_title # 代码块独立成 chunk code_text_parts = [] non_code_parts = [] in_code = False current_block = [] for line in sec_lines: stripped = line.strip() if re.match(r'^```\w*$', stripped): if not in_code: in_code = True current_block = [line] else: current_block.append(line) code_text = '\n'.join(current_block) if len(code_text.strip()) > 20: code_text_parts.append(code_text) current_block = [] in_code = False continue if in_code: current_block.append(line) continue non_code_parts.append(line) # 代码块 chunk for code_text in code_text_parts: chunks.append({ "id": f"{rel_path}#{chunk_id}", "text": f"[代码块] {title_path}\n{code_text}", "title": title_path, "doc_title": doc_title, "source": rel_path, "url": doc_url, "heading": sec["heading"], }) chunk_id += 1 # 非代码部分 chunk non_code_text = '\n'.join(non_code_parts).strip() if not non_code_text or len(non_code_text) < 15: continue if len(non_code_text) <= CHUNK_SIZE: chunks.append({ "id": f"{rel_path}#{chunk_id}", "text": f"{title_path}\n{non_code_text}", "title": title_path, "doc_title": doc_title, "source": rel_path, "url": doc_url, "heading": sec["heading"], }) chunk_id += 1 else: # 按段落拆分大文本 paragraphs = non_code_text.split('\n\n') current = "" for para in paragraphs: para = para.strip() if not para: continue if len(current) + len(para) + 2 <= CHUNK_SIZE: current += ("\n\n" if current else "") + para else: if current: chunks.append({ "id": f"{rel_path}#{chunk_id}", "text": f"{title_path}\n{current}", "title": title_path, "doc_title": doc_title, "source": rel_path, "url": doc_url, "heading": sec["heading"], }) chunk_id += 1 if len(para) > CHUNK_SIZE: # 按行拆 sub_lines = para.split('\n') current = "" for sl in sub_lines: if len(current) + len(sl) + 1 <= CHUNK_SIZE: current += ("\n" if current else "") + sl else: if current: chunks.append({ "id": f"{rel_path}#{chunk_id}", "text": f"{title_path}\n{current}", "title": title_path, "doc_title": doc_title, "source": rel_path, "url": doc_url, "heading": sec["heading"], }) chunk_id += 1 current = sl else: current = para if current.strip(): chunks.append({ "id": f"{rel_path}#{chunk_id}", "text": f"{title_path}\n{current}", "title": title_path, "doc_title": doc_title, "source": rel_path, "url": doc_url, "heading": sec["heading"], }) chunk_id += 1 return chunks # ============ 文档加载 ============ def load_all_docs() -> list[dict]: docs = [] target_dirs = ("收割机", "通用教程") for md_file in sorted(DOCS_SOURCE.rglob("*.md")): rel_path = str(md_file.relative_to(DOCS_SOURCE)) parts = md_file.relative_to(DOCS_SOURCE).parts if parts[0] not in target_dirs: continue if md_file.name in SKIP_FILES: continue try: content = md_file.read_text(encoding="utf-8") doc_title = "" title_match = re.search(r'^title:\s*(.+)$', content, re.MULTILINE) if title_match: doc_title = title_match.group(1).strip().strip('"').strip("'") doc_title = re.sub(r'^\d+\.\s*', '', doc_title) body = clean_markdown(content) if len(body) < 30: continue if not doc_title: doc_title = md_file.stem doc_url = doc_path_to_url(rel_path) chunks = split_document(body, rel_path, doc_title, doc_url) docs.extend(chunks) code_count = sum(1 for c in chunks if "[代码块]" in c["text"]) print(f" 📄 {doc_title}: {len(chunks)} chunks ({code_count} 代码块)") except Exception as e: print(f"⚠️ 跳过 {md_file}: {e}") return docs # ============ 索引管理 ============ def rebuild_index(): INDEX_DIR.mkdir(parents=True, exist_ok=True) docs = load_all_docs() if not docs: print("❌ 没有找到文档") return print(f"\n📚 加载了 {len(docs)} 个文档块") texts = [d["text"] for d in docs] print("🔄 正在生成中文向量...") embeddings = embed_texts(texts) import numpy as np np.save(INDEX_DIR / "embeddings.npy", embeddings) (INDEX_DIR / "metadata.json").write_text( json.dumps(docs, ensure_ascii=False, indent=2), encoding="utf-8" ) save_qa_cache({}) print(f"\n🎉 知识库索引完成，共 {len(docs)} 个文档块") def load_metadata() -> list[dict]: """只加载轻量 metadata，运行时 lite 检索使用。""" meta_path = INDEX_DIR / "metadata.json" if not meta_path.exists(): return [] return json.loads(meta_path.read_text(encoding="utf-8")) def load_index(): """加载向量索引。仅 vector 模式/离线构建环境需要 numpy。""" import numpy as np emb_path = INDEX_DIR / "embeddings.npy" meta_path = INDEX_DIR / "metadata.json" if not emb_path.exists() or not meta_path.exists(): return np.array([]), [] embeddings = np.load(emb_path) metadata = json.loads(meta_path.read_text(encoding="utf-8")) return embeddings, metadata # ============ 问答缓存 ============ def _cache_key(question: str) -> str: normalized = re.sub(r'\s+', '', question.strip().lower()) return hashlib.md5(f"{CACHE_VERSION}:{normalized}".encode()).hexdigest() def load_qa_cache() -> dict: if CACHE_FILE.exists(): try: content = CACHE_FILE.read_text(encoding="utf-8") if content.strip(): return json.loads(content) except Exception: return {} return {} def save_qa_cache(cache: dict): CACHE_FILE.parent.mkdir(parents=True, exist_ok=True) CACHE_FILE.write_text(json.dumps(cache, ensure_ascii=False, indent=2), encoding="utf-8") def get_cached_answer(question: str) -> str | None: cache = load_qa_cache() entry = cache.get(_cache_key(question)) if entry: return entry.get("answer") return None def set_cached_answer(question: str, answer: str): cache = load_qa_cache() cache[_cache_key(question)] = { "question": question, "answer": answer, } if len(cache) > 500: oldest = min(cache.keys(), key=lambda k: cache[k].get("question", "")) del cache[oldest] save_qa_cache(cache) # ============ 关键词→文档映射 ============ KEYWORD_DOC_MAP = { "安装": ["收割机/install.md"], "部署": ["收割机/install.md", "收割机/使用指南.md"], "docker": ["收割机/install.md", "收割机/使用指南.md"], "compose": ["收割机/install.md", "收割机/compose.md"], "群晖": ["收割机/install.md"], "飞牛": ["收割机/install.md"], "cookie": ["收割机/import.md", "通用教程/CookieCloud.md"], "导入": ["收割机/import.md"], "站点导入": ["收割机/import.md"], "签到": ["收割机/common.md"], "通知": ["收割机/common.md", "收割机/home.md"], "下载器": ["收割机/common.md", "收割机/home.md"], "app": ["收割机/app.md"], "辅种": ["收割机/common.md"], "主题": ["收割机/custom-theme.md"], "自定义": ["收割机/custom-add.md", "收割机/custom-theme.md"], "常见问题": ["收割机/common.md"], "微信": ["收割机/wechat.md"], "baidu": ["通用教程/baidu-ocr.md"], "ocr": ["通用教程/baidu-ocr.md"], "frp": ["通用教程/tunnel.md"], "内网穿透": ["通用教程/tunnel.md"], "tailscale": ["通用教程/tailscale-nat.md"], "telegram": ["通用教程/telegram-bot.md"], "阿里云": ["通用教程/aliyun.md"], "emby": ["收割机/common.md"], "做种": ["收割机/common.md"], "日志": ["收割机/common.md"], "代理": ["收割机/common.md"], "ua": ["收割机/common.md"], "登录": ["收割机/使用指南.md"], "添加站点": ["收割机/import.md", "收割机/custom-add.md"], "计划任务": ["收割机/home.md", "收割机/使用指南.md"], "辅种助手": ["收割机/import.md"], "种子迁移": ["收割机/使用指南.md"], "数据库": ["收割机/install.md"], "postgresql": ["收割机/install.md"], "环境变量": ["收割机/install.md"], "webhook": ["收割机/common.md"], "授权": ["收割机/common.md", "收割机/install.md"], "密码": ["收割机/使用指南.md"], "token": ["收割机/使用指南.md"], } def keyword_boost(query: str) -> list[str]: query_lower = query.lower() matched = [] for keyword, docs in KEYWORD_DOC_MAP.items(): if keyword in query_lower: matched.extend(docs) return list(dict.fromkeys(matched)) # ============ 混合检索 ============ def tokenize_query(query: str) -> list[str]: """轻量中文/英文分词：无需 jieba，运行镜像零重依赖。""" q = query.lower().strip() tokens: list[str] = [] # 英文、数字、版本号等 tokens.extend(re.findall(r"[a-z0-9][a-z0-9_.:-]{1,}", q)) # 连续中文词 + 2~4 字滑窗，兼顾“怎么安装/反代/环境变量”等短问法 for block in re.findall(r"[\u4e00-\u9fff]+", q): if len(block) <= 4: tokens.append(block) else: tokens.append(block) for n in (2, 3, 4): tokens.extend(block[i:i + n] for i in range(0, len(block) - n + 1)) # 去重并过滤太泛的单字 return list(dict.fromkeys(t for t in tokens if len(t) >= 2)) def lite_search(query: str, top_k: int = TOP_K) -> list[dict]: """轻量关键词/BM25-like 检索。 运行时不加载 sentence-transformers/torch，大幅降低 Docker 镜像体积。 对当前 200~500 个文档块规模，关键词+领域词表加权的性价比最高。 """ metadata = load_metadata() if not metadata: return [] query_lower = query.lower().strip() tokens = tokenize_query(query) boosted_sources = set(keyword_boost(query)) scored: list[tuple[float, int]] = [] for i, doc in enumerate(metadata): title = (doc.get("title") or "").lower() source = (doc.get("source") or "").lower() text = (doc.get("text") or "").lower() hay = f"{title}\n{source}\n{text}" score = 0.0 # 领域词表强加权 if doc.get("source") in boosted_sources: score += 20.0 # 完整问题/短语命中 if query_lower and query_lower in hay: score += 6.0 if query_lower and query_lower in title: score += 5.0 # token 频次 + 字段权重 for token in tokens: if token in title: score += 4.0 if token in source: score += 3.0 count = text.count(token) if count: # 类 BM25：频次递减，避免长文刷分 score += min(3.0, 1.0 + count ** 0.5) # 配置/安装类问题：代码块和关键配置变量优先，避免模型拿不到原文示例后自行编造 config_query = any(k in query_lower for k in ["compose", "docker", "postgres", "postgresql", "安装", "部署", "配置", "环境变量"]) if config_query and "[代码块]" in text: score += 10.0 config_hits = sum(1 for k in [ "image:", "services:", "environment:", "volumes:", "ports:", "postgres_db", "postgres_user", "postgres_password", "email", "token", "go_web_port", "database", "db_", ] if k in text) score += min(20.0, config_hits * 3.0) # 查询中英文/数字 token 对安装、docker、postgresql 等问题很关键 exact_tokens = re.findall(r"[a-z0-9][a-z0-9_.:-]{1,}", query_lower) for token in exact_tokens: if token in hay: score += 2.5 if score > 0: scored.append((score, i)) scored.sort(reverse=True, key=lambda x: x[0]) docs = [] for score, idx in scored: doc = metadata[idx].copy() doc["score"] = float(score) src_count = sum(1 for d in docs if d["source"] == doc["source"]) if src_count < 2: docs.append(doc) if len(docs) >= top_k: break return docs def vector_search(query: str, top_k: int = TOP_K) -> list[dict]: """向量检索：仅开发/宿主机模式使用，需要 sentence-transformers。""" embeddings, metadata = load_index() if len(embeddings) == 0: return [] query_vec = get_model().encode([query], normalize_embeddings=True) scores = (embeddings @ query_vec.T).flatten() boosted_sources = keyword_boost(query) if boosted_sources: for i, doc in enumerate(metadata): if doc["source"] in boosted_sources: scores[i] += 0.3 query_words = set(re.findall(r'[\u4e00-\u9fff]+', query)) for i, doc in enumerate(metadata): text_lower = doc["text"].lower() matches = sum(1 for w in query_words if w in text_lower) if matches >= 2: scores[i] += 0.15 * matches top_indices = np.argsort(scores)[::-1][:top_k] docs = [] for idx in top_indices: doc = metadata[idx].copy() doc["score"] = float(scores[idx]) src_count = sum(1 for d in docs if d["source"] == doc["source"]) if src_count < 2: docs.append(doc) return docs def expand_context(docs: list[dict], window: int = 1, max_chars: int = 12000, max_sources: int = 5) -> list[dict]: """为命中的 chunk 补充同源相邻上下文，减少回答过短/断章取义。""" if not docs: return [] metadata = load_metadata() by_id = {d.get("id"): d for d in metadata} expanded: list[dict] = [] seen: set[str] = set() total_chars = 0 def parse_id(doc_id: str) -> tuple[str, int] | None: if not doc_id or "#" not in doc_id: return None src, idx = doc_id.rsplit("#", 1) if not idx.isdigit(): return None return src, int(idx) allowed_sources = list(dict.fromkeys(d.get("source") for d in docs if d.get("source")))[:max_sources] allowed_sources_set = set(allowed_sources) for doc in docs: if doc.get("source") not in allowed_sources_set: continue parsed = parse_id(doc.get("id", "")) candidates = [doc] if parsed: src, idx = parsed neighbors = [] for offset in range(-window, window + 1): item = by_id.get(f"{src}#{idx + offset}") if item: neighbors.append(item) candidates = neighbors or [doc] for item in candidates: doc_id = item.get("id", "") if doc_id in seen: continue text_len = len(item.get("text", "")) if total_chars + text_len > max_chars and expanded: continue copy = item.copy() copy.setdefault("score", doc.get("score", 0)) expanded.append(copy) seen.add(doc_id) total_chars += text_len return expanded def search(query: str, top_k: int = TOP_K, expand: bool = True, context_window: int = 1) -> list[dict]: mode = os.environ.get("KNOWLEDGE_SEARCH_MODE", "lite").lower() if mode in {"vector", "embedding", "semantic"}: docs = vector_search(query, top_k=top_k) else: docs = lite_search(query, top_k=top_k) return expand_context(docs, window=context_window) if expand else docs def get_stats() -> dict: metadata = load_metadata() cache = load_qa_cache() return { "total_chunks": len(metadata), "cached_answers": len(cache), "site_url": SITE_URL, "search_mode": os.environ.get("KNOWLEDGE_SEARCH_MODE", "lite"), "embedding_model": EMBED_MODEL, "cache_version": CACHE_VERSION, } # ============ CLI ============ if __name__ == "__main__": import sys if len(sys.argv) > 1 and sys.argv[1] == "rebuild": rebuild_index() elif len(sys.argv) > 1 and sys.argv[1] == "search": query = " ".join(sys.argv[2:]) results = search(query) for r in results: print(f"\n📄 {r['title']} ({r['source']}) [分数:{r.get('score', 0):.3f}]") print(f" URL: {r['url']}") print(f" {r['text'][:200]}...") elif len(sys.argv) > 1 and sys.argv[1] == "stats": print(json.dumps(get_stats(), ensure_ascii=False, indent=2)) elif len(sys.argv) > 1 and sys.argv[1] == "cache": cache = load_qa_cache() for k, v in cache.items(): print(f"\n🔑 {k[:8]}... Q: {v['question'][:50]}") else: print("用法:") print(" python -m knowledge rebuild - 重建索引") print(" python -m knowledge search <查询> - 搜索") print(" python -m knowledge stats - 统计信息") print(" python -m knowledge cache - 查看缓存")