num/knowledge/scripts/local_retriever.py

"""
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()