Files

282 lines
9.2 KiB
Python

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Z 盘知识库检索工具 (ZKB Search)
================================
检索 Z 盘 (挂载自 192.168.1.246 的 Obsidian 知识库) 的文章/笔记。
⚠️ 编码铁律: 本脚本必须用 `py -X utf8` 运行, 否则 Windows 控制台 GBK 输出会乱码!
Windows 下请始终: py -X utf8 zkb_search.py <参数>
用法:
py -X utf8 zkb_search.py search <关键词> # 全文搜索 (文件名+内容)
py -X utf8 zkb_search.py title <关键词> # 按标题搜索 (文件名+frontmatter title)
py -X utf8 zkb_search.py index <关键词> # 搜索索引文件 (wechat-article-links.md 等)
py -X utf8 zkb_search.py recent [N] # 最近 N 篇 (默认10)
py -X utf8 zkb_search.py list # 列出知识库结构
py -X utf8 zkb_search.py read <文件路径> # 安全读取文件 (UTF-8)
py -X utf8 zkb_search.py find <文件名片段> # 按文件名查找
示例:
py -X utf8 zkb_search.py search 量化交易进阶
py -X utf8 zkb_search.py title 因子挖掘
py -X utf8 zkb_search.py recent 5
"""
import os
import sys
import glob
import re
import io
# 强制 UTF-8 输出 (防止 Windows GBK 控制台乱码)
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')
# Z 盘根路径
Z_ROOT = "Z:/"
# 索引文件
INDEX_FILES = [
"index.md", # 根索引 (首页)
"wechat-article-links.md", # 微信文章索引 (最重要)
"log.md",
"research_queue.md",
"self_remediation.md",
]
# 主要搜索目录
SEARCH_DIRS = [
"raw/articles", # 原文文章 (微信抓取等)
"wiki/articles", # LLM Wiki 精读
"wiki/concepts", # LLM Wiki 概念
"wiki/categories", # LLM Wiki 分类索引
"knowledge", # 知识库分类
"daily", # 日报/调研/沉淀
]
def eprint(msg):
print(msg, file=sys.stderr, flush=True)
def safe_read(path):
"""安全读取文件 (UTF-8, 不报错)"""
try:
with open(path, encoding='utf-8', errors='replace') as f:
return f.read()
except Exception as e:
return f"<读取失败: {e}>"
def search_content(keyword, max_results=20):
"""全文搜索: 遍历主要目录, 匹配文件名和内容 (限制文件大小防卡顿)"""
results = []
kw = keyword.lower()
for d in SEARCH_DIRS:
base = os.path.join(Z_ROOT, d)
if not os.path.isdir(base):
continue
for f in glob.glob(os.path.join(base, "*.md")):
try:
# 跳过超大文件 (>500KB) 防卡顿
if os.path.getsize(f) > 500 * 1024:
continue
rel = os.path.relpath(f, Z_ROOT).replace("\\", "/")
fname = os.path.basename(f).lower()
# 文件名匹配
if kw in fname:
results.append((rel, "filename", ""))
if len(results) >= max_results:
return results
continue
# 内容匹配 (只读前50KB)
with open(f, encoding='utf-8', errors='replace') as fh:
content = fh.read(50 * 1024)
if kw in content.lower():
lines = content.split('\n')
match_line = ""
for line in lines:
if kw in line.lower():
match_line = line.strip()[:120]
break
results.append((rel, "content", match_line))
if len(results) >= max_results:
return results
except Exception:
continue
return results
def search_title(keyword, max_results=20):
"""按标题搜索: 只匹配文件名 (快, 不读内容)"""
results = []
kw = keyword.lower()
for d in SEARCH_DIRS:
base = os.path.join(Z_ROOT, d)
if not os.path.isdir(base):
continue
# 只 glob 一级目录, 更快
for f in glob.glob(os.path.join(base, "*.md")):
try:
rel = os.path.relpath(f, Z_ROOT).replace("\\", "/")
fname = os.path.basename(f).lower()
if kw in fname:
results.append((rel, "file_title", fname[:100]))
if len(results) >= max_results:
return results
except Exception:
continue
return results
def search_index(keyword, max_results=30):
"""搜索索引文件 (wechat-article-links.md 等)"""
results = []
kw = keyword.lower()
for idx in INDEX_FILES:
path = os.path.join(Z_ROOT, idx)
if not os.path.exists(path):
continue
content = safe_read(path)
lines = content.split('\n')
for i, line in enumerate(lines):
if kw in line.lower():
# 提取标题部分
title_part = line[:150]
results.append((idx, f"line {i+1}", title_part))
if len(results) >= max_results:
return results
return results
def recent(n=10):
"""最近收录的文章: 从文件名日期前缀排序 (快速, 不读每个文件的mtime)"""
files = []
base = os.path.join(Z_ROOT, "raw/articles")
if os.path.isdir(base):
try:
for entry in os.scandir(base):
if entry.name.endswith(".md"):
files.append(entry.name)
except Exception:
pass
# 从文件名提取 YYYY-MM-DD 日期
def sort_key(name):
m = re.match(r'(\d{4}-\d{2}-\d{2})', name)
return m.group(1) if m else "0000-00-00"
files.sort(key=sort_key, reverse=True)
results = []
for fname in files[:n]:
rel = "raw/articles/" + fname
date_part = sort_key(fname)
results.append((rel, date_part, ""))
return results
def list_structure():
"""列出知识库顶层结构"""
print("=== Z 盘知识库结构 ===")
for d in SEARCH_DIRS:
base = os.path.join(Z_ROOT, d)
if os.path.isdir(base):
count = len(glob.glob(os.path.join(base, "*.md")))
print(f" {d}: {count} 个 md 文件")
print(f"\n=== 索引文件 ===")
for idx in INDEX_FILES:
path = os.path.join(Z_ROOT, idx)
if os.path.exists(path):
size = os.path.getsize(path)
print(f" {idx}: {size/1024:.1f} KB")
def cmd_read(path):
"""安全读取指定文件"""
full = path if os.path.isabs(path) else os.path.join(Z_ROOT, path)
if not os.path.exists(full):
print(f"文件不存在: {path}")
return
content = safe_read(full)
print(content[:5000])
def cmd_find(fragment):
"""按文件名片段查找"""
results = []
for d in SEARCH_DIRS:
base = os.path.join(Z_ROOT, d)
if not os.path.isdir(base):
continue
for f in glob.glob(os.path.join(base, "*.md")):
if fragment.lower() in os.path.basename(f).lower():
rel = os.path.relpath(f, Z_ROOT).replace("\\", "/")
results.append(rel)
for r in results[:30]:
print(r)
print(f"\n{len(results)} 个匹配")
def main():
if len(sys.argv) < 2:
print(__doc__)
return
cmd = sys.argv[1]
if cmd == "search" and len(sys.argv) >= 3:
kw = sys.argv[2]
print(f"=== 全文搜索: {kw} ===")
results = search_content(kw)
if not results:
print("无结果")
for rel, match_type, detail in results:
prefix = "[文件名]" if match_type == "filename" else "[内容]"
print(f" {prefix} {rel}")
if detail:
print(f" {detail}")
print(f"\n{len(results)} 条")
elif cmd == "title" and len(sys.argv) >= 3:
kw = sys.argv[2]
print(f"=== 标题搜索: {kw} ===")
results = search_title(kw)
if not results:
print("无结果")
for rel, match_type, detail in results:
prefix = "[文件名]" if match_type == "file_title" else "[title字段]"
print(f" {prefix} {rel}")
if detail:
print(f" {detail}")
print(f"\n{len(results)} 条")
elif cmd == "index" and len(sys.argv) >= 3:
kw = sys.argv[2]
print(f"=== 索引搜索: {kw} ===")
results = search_index(kw)
if not results:
print("无结果")
for idx, loc, line in results:
print(f" [{idx}:{loc}] {line}")
print(f"\n{len(results)} 条")
elif cmd == "recent":
n = int(sys.argv[2]) if len(sys.argv) >= 3 else 10
print(f"=== 最近 {n} 篇 ===")
for rel, ts, _ in recent(n):
print(f" {ts} {rel}")
elif cmd == "list":
list_structure()
elif cmd == "read" and len(sys.argv) >= 3:
cmd_read(sys.argv[2])
elif cmd == "find" and len(sys.argv) >= 3:
cmd_find(sys.argv[2])
else:
print(__doc__)
if __name__ == "__main__":
main()