""" Local Knowledge Base Retriever ============================== Dify 不可用时的本地 RAG 替代方案。 使用 character n-gram + TF 向量化的方式做检索，无需任何外部 ML 依赖（numpy+scipy only）。 Usage: from local_retriever import LocalRetriever retriever = LocalRetriever() retriever.index_documents() results = retriever.query("主性格数字1的性格特点", top_k=5) """ import os import re import json import hashlib from pathlib import Path from typing import List, Dict, Optional, Tuple import numpy as np # ─── 配置 ──────────────────────────────────────────────────────────────────── KNOWLEDGE_DIR = Path(__file__).resolve().parent.parent / "research" NGRAM_N = 2 # character bigram TOP_K_DEFAULT = 6 # ─── 文档分块 ──────────────────────────────────────────────────────────────── def _parse_yaml_frontmatter(text: str) -> Tuple[Dict, str]: """解析 Markdown YAML frontmatter (--- 包裹的元数据)""" meta = {} body = text m = re.match(r'^---\s*\n(.*?)\n---\s*\n(.*)', text, re.DOTALL) if m: yaml_block = m.group(1) body = m.group(2) for line in yaml_block.strip().split('\n'): if ':' in line: k, v = line.split(':', 1) meta[k.strip()] = v.strip().strip('"').strip("'") return meta, body def _split_sections(text: str, source_file: str) -> List[Dict]: """将文档按 ## 标题拆分为 chunks，每个 chunk 包含上下文""" chunks = [] lines = text.split('\n') current_heading = "前言" current_content = [] heading_level = 0 for line in lines: heading_match = re.match(r'^(#{1,4})\s+(.+)$', line) if heading_match: if current_content: chunks.append({ "id": hashlib.md5(f"{source_file}:{current_heading}".encode()).hexdigest()[:12], "source": source_file, "heading": current_heading, "content": '\n'.join(current_content).strip(), "tokens": len(''.join(current_content)) // 2, }) heading_level = len(heading_match.group(1)) current_heading = heading_match.group(2).strip() current_content = [] else: current_content.append(line) if current_content: chunks.append({ "id": hashlib.md5(f"{source_file}:{current_heading}".encode()).hexdigest()[:12], "source": source_file, "heading": current_heading, "content": '\n'.join(current_content).strip(), "tokens": len(''.join(current_content)) // 2, }) return chunks def load_all_documents(kb_dir: Optional[Path] = None) -> List[Dict]: """加载 knowledge/research/ 下所有 .md 文件，返回 chunks 列表""" if kb_dir is None: kb_dir = KNOWLEDGE_DIR all_chunks = [] md_files = sorted(kb_dir.glob("*.md")) for fpath in md_files: try: text = fpath.read_text(encoding='utf-8') except Exception as e: print(f" [SKIP] {fpath.name}: {e}") continue meta, body = _parse_yaml_frontmatter(text) chunks = _split_sections(body, fpath.name) # 将 YAML 元信息附加到每个 chunk for ch in chunks: ch["file_path"] = str(fpath) ch["doc_title"] = meta.get("title", fpath.stem) ch["doc_tags"] = meta.get("tags", "") ch["doc_purpose"] = meta.get("purpose", "") all_chunks.extend(chunks) return all_chunks # ─── 向量化 ────────────────────────────────────────────────────────────────── def _char_ngrams(text: str, n: int = NGRAM_N) -> List[str]: """生成 character n-grams""" # 去除非中文/字母/数字的字符，保留语义 cleaned = re.sub(r'[^\u4e00-\u9fff\w]', '', text) if len(cleaned) < n: return [cleaned] return [cleaned[i:i + n] for i in range(len(cleaned) - n + 1)] def _build_vocab(chunks: List[Dict], n: int = NGRAM_N) -> Dict[str, int]: """构建 n-gram 词汇表""" vocab = {} for ch in chunks: content = f"{ch.get('doc_title', '')} {ch.get('heading', '')} {ch.get('doc_tags', '')} {ch['content']}" grams = _char_ngrams(content, n) for g in grams: if g not in vocab: vocab[g] = len(vocab) return vocab def _vectorize(text: str, vocab: Dict[str, int], n: int = NGRAM_N) -> np.ndarray: """将文本转为 TF 向量""" vec = np.zeros(len(vocab), dtype=np.float32) grams = _char_ngrams(text, n) for g in grams: idx = vocab.get(g) if idx is not None: vec[idx] += 1.0 # L2 归一化 norm = np.linalg.norm(vec) if norm > 0: vec /= norm return vec # ─── 主检索器 ──────────────────────────────────────────────────────────────── class LocalRetriever: """本地知识库检索器""" def __init__(self, kb_dir: Optional[Path] = None, ngram_n: int = NGRAM_N): self.kb_dir = kb_dir or KNOWLEDGE_DIR self.ngram_n = ngram_n self.chunks: List[Dict] = [] self.vocab: Dict[str, int] = {} self.matrix: Optional[np.ndarray] = None self._indexed = False def index_documents(self) -> Dict: """索引所有文档，返回统计信息""" print(f" Loading documents from: {self.kb_dir}") self.chunks = load_all_documents(self.kb_dir) print(f" Total chunks: {len(self.chunks)}") # 构建词汇表 print(f" Building {self.ngram_n}-gram vocabulary...") self.vocab = _build_vocab(self.chunks, self.ngram_n) print(f" Vocabulary size: {len(self.vocab)}") # 构建文档-词项矩阵 print(" Vectorizing chunks...") rows = [] for i, ch in enumerate(self.chunks): content = f"{ch['doc_title']} {ch['heading']} {ch['doc_tags']} {ch['content']}" vec = _vectorize(content, self.vocab, self.ngram_n) rows.append(vec) if (i + 1) % 50 == 0: print(f" ... {i + 1}/{len(self.chunks)}") self.matrix = np.array(rows, dtype=np.float32) print(f" Matrix shape: {self.matrix.shape}") self._indexed = True return { "documents": len(set(ch["source"] for ch in self.chunks)), "chunks": len(self.chunks), "vocab_size": len(self.vocab), "matrix_shape": list(self.matrix.shape), } def query(self, query_text: str, top_k: int = TOP_K_DEFAULT) -> List[Dict]: """检索与 query 最相关的 chunks""" if not self._indexed: raise RuntimeError("Call index_documents() first") # 向量化 query qvec = _vectorize(query_text, self.vocab, self.ngram_n) # 余弦相似度: cos = dot_product / (norm_q * norm_d) assert self.matrix is not None, "Matrix not initialized" norm_q = np.linalg.norm(qvec) norms_d = np.linalg.norm(self.matrix, axis=1) # 避免除零 denom = norm_q * norms_d denom[denom == 0] = 1e-10 similarities = np.dot(self.matrix, qvec) / denom # 获取 top-k 索引 top_indices = np.argsort(similarities)[::-1][:top_k] results = [] for idx in top_indices: sim = float(similarities[idx]) if sim < 0.01: # 过滤无关结果 break ch = self.chunks[idx] results.append({ "id": ch["id"], "source": ch["source"], "heading": ch["heading"], "doc_title": ch["doc_title"], "content_preview": ch["content"][:200] + ("..." if len(ch["content"]) > 200 else ""), "score": round(sim, 4), "tokens": ch.get("tokens", 0), }) return results def query_formatted(self, query_text: str, top_k: int = TOP_K_DEFAULT) -> str: """返回格式化的检索结果（用于测试报告）""" results = self.query(query_text, top_k) lines = [ f" Query: \"{query_text}\"", f" Results: {len(results)}", ] for i, r in enumerate(results): lines.append(f" [{i + 1}] score={r['score']:.4f} | {r['source']} > {r['heading']}") lines.append(f" {r['content_preview'][:120]}") return '\n'.join(lines) def get_chunks_by_source(self, source: str) -> List[Dict]: """按源文件获取 chunks""" return [ch for ch in self.chunks if ch["source"] == source] def list_sources(self) -> List[str]: """列出所有源文件""" return sorted(set(ch["source"] for ch in self.chunks)) # ─── CLI ───────────────────────────────────────────────────────────────────── def main(): """CLI 入口：索引文档并交互式查询""" retriever = LocalRetriever() stats = retriever.index_documents() print(f"\n{'=' * 50}") print(f" 文档数: {stats['documents']}") print(f" Chunks: {stats['chunks']}") print(f" 词汇量: {stats['vocab_size']}") print(f" 矩阵: {stats['matrix_shape'][0]}x{stats['matrix_shape'][1]}") print(f"{'=' * 50}") test_queries = [ "主性格数字1的性格特点", "财富能量数字组合", "感情婚姻运势分析", "健康疾病数字能量", "五区三组生命阶段", "三角命盘24个位置含义", "组合对判读规则优先级", "事业发展规划建议", ] print("\n --- 测试查询 ---") for q in test_queries: print() print(retriever.query_formatted(q, top_k=3)) print(f"\n{'=' * 50}") print(" Source files indexed:") for src in retriever.list_sources(): n = len(retriever.get_chunks_by_source(src)) print(f" {src} ({n} chunks)") if __name__ == "__main__": main()