Blame for scripts/hybrid_search.py

from __future__ import annotations
# llm-wiki-version: 2.0.0
# runtime: dev-only

"""
hybrid_search.py — v2 混合检索 (BM25 + 图遍历).

Rohit G00 v2 范式：BM25 (keyword) + 向量 (semantic) + 图遍历 (structural).
v2.0 实现: BM25 over wiki 页面 + 图遍历 (follows supersedes/contradicts links).
v2.1+ 可加 embedding-based vector search.

Usage:
  python3 scripts/hybrid_search.py "货币主权"
  python3 scripts/hybrid_search.py "DIAP" --depth 2
  python3 scripts/hybrid_search.py "价值函数" --json --top 5
"""

import argparse
import json
import math
import re
import sys
from collections import Counter, defaultdict
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
WIKI_ROOT = ROOT / "docs" / "wiki"

FRONTMATTER_RE = re.compile(r"^---\n(.*?)\n---", re.DOTALL)
TOKEN_RE = re.compile(r"[\w一-鿿]+", re.UNICODE)  # 英文/数字/汉字

# BM25 parameters
K1 = 1.5
B = 0.75


def parse_frontmatter(text: str) -> dict[str, str]:
    m = FRONTMATTER_RE.match(text)
    if not m:
        return {}
    out: dict[str, str] = {}
    for line in m.group(1).splitlines():
        if ":" in line:
            key, _, value = line.partition(":")
            out[key.strip()] = value.strip()
    return out


def tokenize(text: str) -> list[str]:
    """CJK + ASCII tokenization. Each CJK char is its own token, ASCII words grouped."""
    tokens: list[str] = []
    for m in TOKEN_RE.finditer(text):
        s = m.group(0)
        if re.search(r"[一-鿿]", s):
            tokens.extend(s)  # CJK: one char per token
        else:
            tokens.append(s.lower())
    return tokens


def bm25_score(query_tokens: list[str], doc_tokens: list[str], avg_dl: float, df: dict[str, int], n_docs: int) -> float:
    """Compute BM25 score for a single document."""
    score = 0.0
    doc_len = len(doc_tokens)
    tf = Counter(doc_tokens)
    for q in query_tokens:
        if q not in tf:
            continue
        idf = math.log((n_docs - df[q] + 0.5) / (df[q] + 0.5) + 1)
        denom = tf[q] + K1 * (1 - B + B * doc_len / avg_dl)
        score += idf * (tf[q] * (K1 + 1)) / denom
    return score


def main() -> int:
    parser = argparse.ArgumentParser(description="v2 hybrid search (BM25 + graph)")
    parser.add_argument("query", help="search query")
    parser.add_argument("--wiki-root", type=Path, default=WIKI_ROOT)
    parser.add_argument("--depth", type=int, default=1, help="graph traversal depth (0=no graph)")
    parser.add_argument("--top", type=int, default=5, help="top N results")
    parser.add_argument("--json", action="store_true", help="emit JSON output")
    args = parser.parse_args()

    if not args.wiki_root.exists():
        print(f"hybrid_search: {args.wiki_root} does not exist")
        return 1

    md_files = sorted(args.wiki_root.rglob("*.md"))
    docs: list[dict] = []
    for path in md_files:
        text = path.read_text(encoding="utf-8")
        # Strip frontmatter for tokenization
        body = FRONTMATTER_RE.sub("", text, count=1)
        tokens = tokenize(body)
        fm = parse_frontmatter(text)
        rel = path.relative_to(args.wiki_root).as_posix()
        docs.append({"rel": rel, "tokens": tokens, "fm": fm, "path": path})

    if not docs:
        print("hybrid_search: no wiki pages found")
        return 1

    # Build DF
    df: dict[str, int] = defaultdict(int)
    for d in docs:
        for term in set(d["tokens"]):
            df[term] += 1

    avg_dl = sum(len(d["tokens"]) for d in docs) / len(docs)
    query_tokens = tokenize(args.query)

    # BM25 scores
    bm25_hits: list[tuple[float, dict]] = []
    for d in docs:
        score = bm25_score(query_tokens, d["tokens"], avg_dl, df, len(docs))
        if score > 0:
            bm25_hits.append((score, d))
    bm25_hits.sort(key=lambda x: -x[0])

    # Build graph from supersedes + contradicts
    graph: dict[str, set[str]] = defaultdict(set)
    for d in docs:
        rel = d["rel"]
        for edge_field in ("supersedes", "contradicts"):
            raw = d["fm"].get(edge_field, "")
            if not raw:
                continue
            for t in re.findall(r"\./[^\s,\]]+", raw):
                graph[rel].add(t.lstrip("./"))

    # Graph traversal: for each top BM25 hit, expand neighbors up to depth
    bm25_paths = {d["rel"]: score for score, d in bm25_hits[:args.top]}
    expanded_paths: dict[str, list[str]] = {}

    if args.depth > 0:
        # BFS from top hits
        visited: set[str] = set()
        queue: list[tuple[str, int, list[str]]] = []
        for rel in bm25_paths:
            queue.append((rel, 0, [rel]))
            visited.add(rel)
        while queue:
            node, d, path = queue.pop(0)
            if d >= args.depth:
                continue
            for nxt in graph.get(node, set()):
                if nxt not in visited and nxt in {d2["rel"] for d2 in docs}:
                    visited.add(nxt)
                    new_path = path + [nxt]
                    queue.append((nxt, d + 1, new_path))
                    if d + 1 == args.depth:
                        # Add as expansion
                        expanded_paths.setdefault(nxt, []).extend(new_path)
        # Flatten
        graph_hits: list[tuple[str, int, list[str]]] = []
        for nxt, paths in expanded_paths.items():
            for p in paths:
                graph_hits.append((0.5 / (len(p)), d, p))
        graph_hits.sort(key=lambda x: -x[0])
    else:
        graph_hits = []

    # Reciprocal Rank Fusion (RRF)
    rrf_scores: dict[str, float] = defaultdict(float)
    for rank, (score, d) in enumerate(bm25_hits, start=1):
        rrf_scores[d["rel"]] += 1.0 / (60 + rank)
    for rank, (score, _, path) in enumerate(graph_hits, start=1):
        rrf_scores[path[-1]] += 0.5 / (60 + rank)

    # Final ranking
    final_ranked = sorted(rrf_scores.items(), key=lambda kv: -kv[1])[:args.top]

    # Build snippets
    results: list[dict] = []
    for rel, rrf in final_ranked:
        doc = next(d for d in docs if d["rel"] == rel)
        body = FRONTMATTER_RE.sub("", doc["path"].read_text(encoding="utf-8"), count=1)
        # Find first matching line
        snippet = ""
        for line in body.splitlines():
            if any(q.lower() in line.lower() for q in [args.query]):
                snippet = line.strip()[:200]
                break
        if not snippet:
            snippet = body.splitlines()[0][:200] if body.splitlines() else ""
        results.append({
            "rel": rel,
            "rrf_score": round(rrf, 5),
            "title": doc["fm"].get("title", ""),
            "stage": doc["fm"].get("stage", ""),
            "snippet": snippet,
        })

    report = {
        "query": args.query,
        "query_tokens": query_tokens,
        "wiki_pages_scanned": len(docs),
        "bm25_hits": len(bm25_hits),
        "graph_depth": args.depth,
        "results": results,
    }

    if args.json:
        print(json.dumps(report, indent=2, ensure_ascii=False))
    else:
        print(f"hybrid_search: query={args.query!r} (tokens: {query_tokens[:10]}{'...' if len(query_tokens) > 10 else ''})")
        print(f"- scanned: {len(docs)} pages")
        print(f"- BM25 hits: {len(bm25_hits)}")
        print(f"- graph depth: {args.depth}")
        print(f"- top {args.top} results (RRF):")
        for r in results:
            print(f"  {r['rrf_score']:.4f}  {r['rel']}  [{r['stage']}]  {r['title']}")
            if r['snippet']:
                print(f"           {r['snippet'][:120]}")

    return 0


if __name__ == "__main__":
    sys.exit(main())