150 lines
5.0 KiB
Python
150 lines
5.0 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
graph-search.py — 图感知搜索
|
||
|
||
流程：
|
||
1. qmd vsearch 获取语义匹配结果（Top 10）
|
||
2. 对每个结果，扩展相邻节点：
|
||
- 默认：读取各页面 relations 字段（实时，慢但零状态）
|
||
- --use-graph：读取预构建的 knowledge-graph.json（快，需先 build）
|
||
3. 去重后返回（直接关联 + 关系扩展）
|
||
|
||
用法：
|
||
python tools/scripts/graph-search.py "查询词" [--count 15] [--depth 1] [--use-graph]
|
||
|
||
依赖：
|
||
- qmd 已安装（通过 node）
|
||
- --use-graph 需先运行 knowledge-graph.py build
|
||
"""
|
||
import argparse
|
||
import json
|
||
import re
|
||
import subprocess
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
try:
|
||
import yaml
|
||
except ImportError:
|
||
yaml = None
|
||
|
||
QMD = r'node "C:\Users\hhhh2024\AppData\Roaming\npm\node_modules\@tobilu\qmd\dist\cli\qmd.js"'
|
||
WIKI = Path(r"D:\Applications\app\kepano-obsidian-main\wiki")
|
||
GRAPH_DB = Path(r"D:\Applications\app\kepano-obsidian-main\tools\data\knowledge-graph.json")
|
||
|
||
|
||
def qmd_search(query: str, count: int = 10) -> list[str]:
|
||
"""调用 qmd vsearch 获取匹配的页面名列表"""
|
||
cmd = f'{QMD} vsearch "{query}" -c wiki -n {count}'
|
||
result = subprocess.run(cmd, capture_output=True, text=True, shell=True, timeout=30)
|
||
lines = [line.strip() for line in result.stdout.splitlines() if line.strip()]
|
||
return lines
|
||
|
||
|
||
def get_relations(page_name: str) -> list[dict]:
|
||
"""读取页面的 relations 字段"""
|
||
fp = WIKI / f"{page_name}.md"
|
||
if not fp.exists():
|
||
return []
|
||
content = fp.read_text(encoding="utf-8")
|
||
m = re.match(r"^---\n(.*?)\n---", content, re.DOTALL)
|
||
if not m:
|
||
return []
|
||
if not yaml:
|
||
return []
|
||
try:
|
||
front = yaml.safe_load(m.group(1)) or {}
|
||
except:
|
||
return []
|
||
return front.get("relations", [])
|
||
|
||
|
||
def expand_nodes(nodes: list[str], depth: int = 1) -> list[str]:
|
||
"""从起始节点出发，沿 relations 扩展相邻节点（实时读取页面 frontmatter）"""
|
||
expanded = list(nodes)
|
||
frontier = list(nodes)
|
||
for _ in range(depth):
|
||
next_frontier = []
|
||
for node in frontier:
|
||
rels = get_relations(node)
|
||
for r in rels:
|
||
target = r.get("target", "").strip("[]")
|
||
if target and target not in expanded:
|
||
expanded.append(target)
|
||
next_frontier.append(target)
|
||
frontier = next_frontier
|
||
if not frontier:
|
||
break
|
||
return expanded
|
||
|
||
|
||
def expand_nodes_via_graph(nodes: list[str], depth: int = 1, min_conf: int = 3) -> list[str]:
|
||
"""通过预构建的 knowledge-graph.json 扩展（更快，无 IO）"""
|
||
if not GRAPH_DB.exists():
|
||
return list(nodes)
|
||
graph = json.loads(GRAPH_DB.read_text(encoding="utf-8"))
|
||
adj: dict[str, list] = {}
|
||
for e in graph["edges"]:
|
||
adj.setdefault(e["source"], []).append((e["target"], e.get("confidence", 3)))
|
||
adj.setdefault(e["target"], []).append((e["source"], e.get("confidence", 3)))
|
||
|
||
expanded = list(nodes)
|
||
frontier = list(nodes)
|
||
for _ in range(depth):
|
||
next_frontier = []
|
||
for node in frontier:
|
||
for target, conf in adj.get(node, []):
|
||
if conf >= min_conf and target not in expanded:
|
||
expanded.append(target)
|
||
next_frontier.append(target)
|
||
frontier = next_frontier
|
||
if not frontier:
|
||
break
|
||
return expanded
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="Graph-aware search for wiki pages")
|
||
parser.add_argument("query", help="Search query")
|
||
parser.add_argument("--count", type=int, default=15, help="Max results")
|
||
parser.add_argument("--depth", type=int, default=1, help="Graph expansion depth")
|
||
parser.add_argument("--use-graph", action="store_true",
|
||
help="Use prebuilt knowledge-graph.json instead of live frontmatter")
|
||
args = parser.parse_args()
|
||
|
||
print(f"=== Searching: {args.query} ===")
|
||
results = qmd_search(args.query, count=max(10, args.count))
|
||
if not results:
|
||
print(" No results from qmd")
|
||
return
|
||
|
||
print(f"\nDirect matches ({len(results)}):")
|
||
for r in results:
|
||
print(f" [[{r}]]")
|
||
|
||
if args.use_graph and not GRAPH_DB.exists():
|
||
print(f"\n (note: {GRAPH_DB.name} not found, run `knowledge-graph.py build`; "
|
||
f"falling back to live frontmatter)")
|
||
|
||
if args.use_graph and GRAPH_DB.exists():
|
||
expanded = expand_nodes_via_graph(results, depth=args.depth)
|
||
source = "knowledge-graph.json"
|
||
else:
|
||
expanded = expand_nodes(results, depth=args.depth)
|
||
source = "live frontmatter"
|
||
|
||
new = [e for e in expanded if e not in results]
|
||
if new:
|
||
print(f"\nGraph-expanded (via {source}, depth={args.depth}):")
|
||
for n in new[:args.count]:
|
||
print(f" [[{n}]]")
|
||
if len(new) > args.count:
|
||
print(f" ... and {len(new) - args.count} more")
|
||
|
||
print(f"\nTotal unique: {len(expanded)}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|