ngfchl
720 lines
24 KiB
Python
720 lines
24 KiB
Python
"""知识库管理模块 - RAG 检索增强生成
|
|
|
|
架构：文档 → 精细分块(代码块独立/标题层级切分) → 中文embedding(BGE) + numpy余弦相似度 → 关键词+向量混合检索 → LLM生成回答
|
|
"""
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import re
|
|
from pathlib import Path
|
|
from urllib.parse import quote
|
|
|
|
# ============ 配置 ============
|
|
|
|
KNOWLEDGE_DIR = Path(__file__).parent
|
|
INDEX_DIR = KNOWLEDGE_DIR / "index"
|
|
CACHE_FILE = KNOWLEDGE_DIR / "qa_cache.json"
|
|
CACHE_VERSION = "v3-detail-grounded"
|
|
DOCS_SOURCE = Path("/vol3/1000/Code/ngfchl.github.io/docs")
|
|
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434")
|
|
OLLAMA_MODEL = os.environ.get("OLLAMA_MODEL", "qwen3:8b")
|
|
SITE_URL = "https://ptools.fun"
|
|
|
|
# 切分参数
|
|
CHUNK_SIZE = 900
|
|
CHUNK_OVERLAP = 120
|
|
TOP_K = 10
|
|
|
|
# 中文 embedding 模型
|
|
EMBED_MODEL = "BAAI/bge-small-zh-v1.5"
|
|
|
|
# 需要跳过的文件
|
|
SKIP_FILES = {"agreement.md", "index.md", "password.md", "termnal.md", "ssh.md"}
|
|
SKIP_SECTION_KEYWORDS = ["油猴", "tampermonkey", "greasemonkey", "userscript"]
|
|
|
|
|
|
# ============ Embedding ============
|
|
|
|
_model = None
|
|
|
|
|
|
def get_model():
|
|
global _model
|
|
if _model is None:
|
|
from sentence_transformers import SentenceTransformer
|
|
cache_dir = str(KNOWLEDGE_DIR / "models")
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = ""
|
|
print(f"📦 加载中文 embedding 模型: {EMBED_MODEL}")
|
|
_model = SentenceTransformer(EMBED_MODEL, cache_folder=cache_dir)
|
|
print("✅ 模型加载完成")
|
|
return _model
|
|
|
|
|
|
def embed_texts(texts: list[str]):
|
|
model = get_model()
|
|
return model.encode(texts, normalize_embeddings=True, show_progress_bar=True,
|
|
batch_size=32, convert_to_numpy=True)
|
|
|
|
|
|
# ============ URL 生成 ============
|
|
|
|
def doc_path_to_url(rel_path: str) -> str:
|
|
html_path = rel_path.replace(".md", ".html")
|
|
parts = html_path.split("/")
|
|
encoded = "/".join(quote(p, safe="") for p in parts)
|
|
return f"{SITE_URL}/{encoded}"
|
|
|
|
|
|
# ============ 文档清洗（保留结构） ============
|
|
|
|
def clean_markdown(text: str) -> str:
|
|
"""清洗但保留代码块标记和链接"""
|
|
text = re.sub(r'<[^>]+>', '', text)
|
|
# 图片替换为 alt 文本
|
|
text = re.sub(r'!\[([^\]]*)\]\([^\)]+\)', r'[图片:\1]', text)
|
|
# 链接保留文本+URL
|
|
text = re.sub(r'\[([^\]]+)\]\(([^\)]+)\)', r'\1 (\2)', text)
|
|
# frontmatter
|
|
text = re.sub(r'^---\n.*?---\n', '', text, flags=re.DOTALL)
|
|
text = re.sub(r'\n{3,}', '\n\n', text)
|
|
return text.strip()
|
|
|
|
|
|
# ============ 代码块检测 ============
|
|
|
|
def find_code_blocks(text: str) -> list[tuple[int, int]]:
|
|
"""找到所有代码块的起止位置"""
|
|
blocks = []
|
|
lines = text.split('\n')
|
|
in_block = False
|
|
start = -1
|
|
|
|
for i, line in enumerate(lines):
|
|
stripped = line.strip()
|
|
if re.match(r'^```\w*$', stripped):
|
|
if not in_block:
|
|
in_block = True
|
|
start = i
|
|
else:
|
|
blocks.append((start, i))
|
|
in_block = False
|
|
return blocks
|
|
|
|
|
|
# ============ 精细分块 ============
|
|
|
|
def split_document(text: str, rel_path: str, doc_title: str, doc_url: str) -> list[dict]:
|
|
"""
|
|
精细分块策略：
|
|
1. 代码块作为独立 chunk（不拆分）
|
|
2. 每个标题层级下的内容独立成 chunk
|
|
3. 大段落按 CHUNK_SIZE 拆分
|
|
4. 每个 chunk 带上标题路径上下文
|
|
"""
|
|
chunks = []
|
|
lines = text.split('\n')
|
|
|
|
# 先找代码块位置
|
|
code_blocks = find_code_blocks(text)
|
|
code_block_lines = set()
|
|
for start, end in code_blocks:
|
|
for i in range(start, end + 1):
|
|
code_block_lines.add(i)
|
|
|
|
# 按标题切分
|
|
sections = [] # [(level, heading, line_start, line_end)]
|
|
current_heading_stack = [] # [(level, heading)]
|
|
current_start = 0
|
|
|
|
for i, line in enumerate(lines):
|
|
if i in code_block_lines:
|
|
continue
|
|
m = re.match(r'^(#{1,6})\s+(.+)', line.strip())
|
|
if m:
|
|
level = len(m.group(1))
|
|
heading = m.group(2).strip()
|
|
# 保存之前的 section
|
|
if i > current_start:
|
|
sections.append({
|
|
"level": current_heading_stack[-1][0] if current_heading_stack else 0,
|
|
"heading": current_heading_stack[-1][1] if current_heading_stack else "",
|
|
"path": [h for _, h in current_heading_stack],
|
|
"line_start": current_start,
|
|
"line_end": i - 1,
|
|
})
|
|
# 更新 heading stack
|
|
while current_heading_stack and current_heading_stack[-1][0] >= level:
|
|
current_heading_stack.pop()
|
|
current_heading_stack.append((level, heading))
|
|
current_start = i
|
|
|
|
# 最后一个 section
|
|
sections.append({
|
|
"level": current_heading_stack[-1][0] if current_heading_stack else 0,
|
|
"heading": current_heading_stack[-1][1] if current_heading_stack else "",
|
|
"path": [h for _, h in current_heading_stack],
|
|
"line_start": current_start,
|
|
"line_end": len(lines) - 1,
|
|
})
|
|
|
|
chunk_id = 0
|
|
|
|
for sec in sections:
|
|
sec_lines = lines[sec["line_start"]:sec["line_end"] + 1]
|
|
sec_text = '\n'.join(sec_lines).strip()
|
|
|
|
if len(sec_text) < 15:
|
|
continue
|
|
|
|
# 检查是否包含油猴关键词
|
|
if any(kw in sec_text.lower() for kw in SKIP_SECTION_KEYWORDS):
|
|
continue
|
|
|
|
# 构建标题路径
|
|
title_path = " > ".join(sec["path"]) if sec["path"] else doc_title
|
|
if title_path and title_path != doc_title:
|
|
title_path = f"{doc_title} > {title_path}"
|
|
else:
|
|
title_path = doc_title
|
|
|
|
# 代码块独立成 chunk
|
|
code_text_parts = []
|
|
non_code_parts = []
|
|
in_code = False
|
|
current_block = []
|
|
|
|
for line in sec_lines:
|
|
stripped = line.strip()
|
|
if re.match(r'^```\w*$', stripped):
|
|
if not in_code:
|
|
in_code = True
|
|
current_block = [line]
|
|
else:
|
|
current_block.append(line)
|
|
code_text = '\n'.join(current_block)
|
|
if len(code_text.strip()) > 20:
|
|
code_text_parts.append(code_text)
|
|
current_block = []
|
|
in_code = False
|
|
continue
|
|
|
|
if in_code:
|
|
current_block.append(line)
|
|
continue
|
|
|
|
non_code_parts.append(line)
|
|
|
|
# 代码块 chunk
|
|
for code_text in code_text_parts:
|
|
chunks.append({
|
|
"id": f"{rel_path}#{chunk_id}",
|
|
"text": f"[代码块] {title_path}\n{code_text}",
|
|
"title": title_path,
|
|
"doc_title": doc_title,
|
|
"source": rel_path,
|
|
"url": doc_url,
|
|
"heading": sec["heading"],
|
|
})
|
|
chunk_id += 1
|
|
|
|
# 非代码部分 chunk
|
|
non_code_text = '\n'.join(non_code_parts).strip()
|
|
if not non_code_text or len(non_code_text) < 15:
|
|
continue
|
|
|
|
if len(non_code_text) <= CHUNK_SIZE:
|
|
chunks.append({
|
|
"id": f"{rel_path}#{chunk_id}",
|
|
"text": f"{title_path}\n{non_code_text}",
|
|
"title": title_path,
|
|
"doc_title": doc_title,
|
|
"source": rel_path,
|
|
"url": doc_url,
|
|
"heading": sec["heading"],
|
|
})
|
|
chunk_id += 1
|
|
else:
|
|
# 按段落拆分大文本
|
|
paragraphs = non_code_text.split('\n\n')
|
|
current = ""
|
|
for para in paragraphs:
|
|
para = para.strip()
|
|
if not para:
|
|
continue
|
|
if len(current) + len(para) + 2 <= CHUNK_SIZE:
|
|
current += ("\n\n" if current else "") + para
|
|
else:
|
|
if current:
|
|
chunks.append({
|
|
"id": f"{rel_path}#{chunk_id}",
|
|
"text": f"{title_path}\n{current}",
|
|
"title": title_path,
|
|
"doc_title": doc_title,
|
|
"source": rel_path,
|
|
"url": doc_url,
|
|
"heading": sec["heading"],
|
|
})
|
|
chunk_id += 1
|
|
if len(para) > CHUNK_SIZE:
|
|
# 按行拆
|
|
sub_lines = para.split('\n')
|
|
current = ""
|
|
for sl in sub_lines:
|
|
if len(current) + len(sl) + 1 <= CHUNK_SIZE:
|
|
current += ("\n" if current else "") + sl
|
|
else:
|
|
if current:
|
|
chunks.append({
|
|
"id": f"{rel_path}#{chunk_id}",
|
|
"text": f"{title_path}\n{current}",
|
|
"title": title_path,
|
|
"doc_title": doc_title,
|
|
"source": rel_path,
|
|
"url": doc_url,
|
|
"heading": sec["heading"],
|
|
})
|
|
chunk_id += 1
|
|
current = sl
|
|
else:
|
|
current = para
|
|
if current.strip():
|
|
chunks.append({
|
|
"id": f"{rel_path}#{chunk_id}",
|
|
"text": f"{title_path}\n{current}",
|
|
"title": title_path,
|
|
"doc_title": doc_title,
|
|
"source": rel_path,
|
|
"url": doc_url,
|
|
"heading": sec["heading"],
|
|
})
|
|
chunk_id += 1
|
|
|
|
return chunks
|
|
|
|
|
|
# ============ 文档加载 ============
|
|
|
|
def load_all_docs() -> list[dict]:
|
|
docs = []
|
|
target_dirs = ("收割机", "通用教程")
|
|
|
|
for md_file in sorted(DOCS_SOURCE.rglob("*.md")):
|
|
rel_path = str(md_file.relative_to(DOCS_SOURCE))
|
|
parts = md_file.relative_to(DOCS_SOURCE).parts
|
|
|
|
if parts[0] not in target_dirs:
|
|
continue
|
|
if md_file.name in SKIP_FILES:
|
|
continue
|
|
|
|
try:
|
|
content = md_file.read_text(encoding="utf-8")
|
|
doc_title = ""
|
|
title_match = re.search(r'^title:\s*(.+)$', content, re.MULTILINE)
|
|
if title_match:
|
|
doc_title = title_match.group(1).strip().strip('"').strip("'")
|
|
doc_title = re.sub(r'^\d+\.\s*', '', doc_title)
|
|
|
|
body = clean_markdown(content)
|
|
if len(body) < 30:
|
|
continue
|
|
if not doc_title:
|
|
doc_title = md_file.stem
|
|
|
|
doc_url = doc_path_to_url(rel_path)
|
|
chunks = split_document(body, rel_path, doc_title, doc_url)
|
|
docs.extend(chunks)
|
|
|
|
code_count = sum(1 for c in chunks if "[代码块]" in c["text"])
|
|
print(f" 📄 {doc_title}: {len(chunks)} chunks ({code_count} 代码块)")
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ 跳过 {md_file}: {e}")
|
|
|
|
return docs
|
|
|
|
|
|
# ============ 索引管理 ============
|
|
|
|
def rebuild_index():
|
|
INDEX_DIR.mkdir(parents=True, exist_ok=True)
|
|
docs = load_all_docs()
|
|
if not docs:
|
|
print("❌ 没有找到文档")
|
|
return
|
|
|
|
print(f"\n📚 加载了 {len(docs)} 个文档块")
|
|
|
|
texts = [d["text"] for d in docs]
|
|
print("🔄 正在生成中文向量...")
|
|
embeddings = embed_texts(texts)
|
|
|
|
import numpy as np
|
|
np.save(INDEX_DIR / "embeddings.npy", embeddings)
|
|
(INDEX_DIR / "metadata.json").write_text(
|
|
json.dumps(docs, ensure_ascii=False, indent=2), encoding="utf-8"
|
|
)
|
|
|
|
save_qa_cache({})
|
|
print(f"\n🎉 知识库索引完成，共 {len(docs)} 个文档块")
|
|
|
|
|
|
def load_metadata() -> list[dict]:
|
|
"""只加载轻量 metadata，运行时 lite 检索使用。"""
|
|
meta_path = INDEX_DIR / "metadata.json"
|
|
if not meta_path.exists():
|
|
return []
|
|
return json.loads(meta_path.read_text(encoding="utf-8"))
|
|
|
|
|
|
def load_index():
|
|
"""加载向量索引。仅 vector 模式/离线构建环境需要 numpy。"""
|
|
import numpy as np
|
|
|
|
emb_path = INDEX_DIR / "embeddings.npy"
|
|
meta_path = INDEX_DIR / "metadata.json"
|
|
if not emb_path.exists() or not meta_path.exists():
|
|
return np.array([]), []
|
|
embeddings = np.load(emb_path)
|
|
metadata = json.loads(meta_path.read_text(encoding="utf-8"))
|
|
return embeddings, metadata
|
|
|
|
|
|
# ============ 问答缓存 ============
|
|
|
|
def _cache_key(question: str) -> str:
|
|
normalized = re.sub(r'\s+', '', question.strip().lower())
|
|
return hashlib.md5(f"{CACHE_VERSION}:{normalized}".encode()).hexdigest()
|
|
|
|
|
|
def load_qa_cache() -> dict:
|
|
if CACHE_FILE.exists():
|
|
try:
|
|
content = CACHE_FILE.read_text(encoding="utf-8")
|
|
if content.strip():
|
|
return json.loads(content)
|
|
except Exception:
|
|
return {}
|
|
return {}
|
|
|
|
|
|
def save_qa_cache(cache: dict):
|
|
CACHE_FILE.parent.mkdir(parents=True, exist_ok=True)
|
|
CACHE_FILE.write_text(json.dumps(cache, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
|
|
|
|
def get_cached_answer(question: str) -> str | None:
|
|
cache = load_qa_cache()
|
|
entry = cache.get(_cache_key(question))
|
|
if entry:
|
|
return entry.get("answer")
|
|
return None
|
|
|
|
|
|
def set_cached_answer(question: str, answer: str):
|
|
cache = load_qa_cache()
|
|
cache[_cache_key(question)] = {
|
|
"question": question,
|
|
"answer": answer,
|
|
}
|
|
if len(cache) > 500:
|
|
oldest = min(cache.keys(), key=lambda k: cache[k].get("question", ""))
|
|
del cache[oldest]
|
|
save_qa_cache(cache)
|
|
|
|
|
|
# ============ 关键词→文档映射 ============
|
|
|
|
KEYWORD_DOC_MAP = {
|
|
"安装": ["收割机/install.md"],
|
|
"部署": ["收割机/install.md", "收割机/使用指南.md"],
|
|
"docker": ["收割机/install.md", "收割机/使用指南.md"],
|
|
"compose": ["收割机/install.md", "收割机/compose.md"],
|
|
"群晖": ["收割机/install.md"],
|
|
"飞牛": ["收割机/install.md"],
|
|
"cookie": ["收割机/import.md", "通用教程/CookieCloud.md"],
|
|
"导入": ["收割机/import.md"],
|
|
"站点导入": ["收割机/import.md"],
|
|
"签到": ["收割机/common.md"],
|
|
"通知": ["收割机/common.md", "收割机/home.md"],
|
|
"下载器": ["收割机/common.md", "收割机/home.md"],
|
|
"app": ["收割机/app.md"],
|
|
"辅种": ["收割机/common.md"],
|
|
"主题": ["收割机/custom-theme.md"],
|
|
"自定义": ["收割机/custom-add.md", "收割机/custom-theme.md"],
|
|
"常见问题": ["收割机/common.md"],
|
|
"微信": ["收割机/wechat.md"],
|
|
"baidu": ["通用教程/baidu-ocr.md"],
|
|
"ocr": ["通用教程/baidu-ocr.md"],
|
|
"frp": ["通用教程/tunnel.md"],
|
|
"内网穿透": ["通用教程/tunnel.md"],
|
|
"tailscale": ["通用教程/tailscale-nat.md"],
|
|
"telegram": ["通用教程/telegram-bot.md"],
|
|
"阿里云": ["通用教程/aliyun.md"],
|
|
"emby": ["收割机/common.md"],
|
|
"做种": ["收割机/common.md"],
|
|
"日志": ["收割机/common.md"],
|
|
"代理": ["收割机/common.md"],
|
|
"ua": ["收割机/common.md"],
|
|
"登录": ["收割机/使用指南.md"],
|
|
"添加站点": ["收割机/import.md", "收割机/custom-add.md"],
|
|
"计划任务": ["收割机/home.md", "收割机/使用指南.md"],
|
|
"辅种助手": ["收割机/import.md"],
|
|
"种子迁移": ["收割机/使用指南.md"],
|
|
"数据库": ["收割机/install.md"],
|
|
"postgresql": ["收割机/install.md"],
|
|
"环境变量": ["收割机/install.md"],
|
|
"webhook": ["收割机/common.md"],
|
|
"授权": ["收割机/common.md", "收割机/install.md"],
|
|
"密码": ["收割机/使用指南.md"],
|
|
"token": ["收割机/使用指南.md"],
|
|
}
|
|
|
|
|
|
def keyword_boost(query: str) -> list[str]:
|
|
query_lower = query.lower()
|
|
matched = []
|
|
for keyword, docs in KEYWORD_DOC_MAP.items():
|
|
if keyword in query_lower:
|
|
matched.extend(docs)
|
|
return list(dict.fromkeys(matched))
|
|
|
|
|
|
# ============ 混合检索 ============
|
|
|
|
def tokenize_query(query: str) -> list[str]:
|
|
"""轻量中文/英文分词：无需 jieba，运行镜像零重依赖。"""
|
|
q = query.lower().strip()
|
|
tokens: list[str] = []
|
|
|
|
# 英文、数字、版本号等
|
|
tokens.extend(re.findall(r"[a-z0-9][a-z0-9_.:-]{1,}", q))
|
|
|
|
# 连续中文词 + 2~4 字滑窗，兼顾“怎么安装/反代/环境变量”等短问法
|
|
for block in re.findall(r"[\u4e00-\u9fff]+", q):
|
|
if len(block) <= 4:
|
|
tokens.append(block)
|
|
else:
|
|
tokens.append(block)
|
|
for n in (2, 3, 4):
|
|
tokens.extend(block[i:i + n] for i in range(0, len(block) - n + 1))
|
|
|
|
# 去重并过滤太泛的单字
|
|
return list(dict.fromkeys(t for t in tokens if len(t) >= 2))
|
|
|
|
|
|
def lite_search(query: str, top_k: int = TOP_K) -> list[dict]:
|
|
"""轻量关键词/BM25-like 检索。
|
|
|
|
运行时不加载 sentence-transformers/torch，大幅降低 Docker 镜像体积。
|
|
对当前 200~500 个文档块规模，关键词+领域词表加权的性价比最高。
|
|
"""
|
|
metadata = load_metadata()
|
|
if not metadata:
|
|
return []
|
|
|
|
query_lower = query.lower().strip()
|
|
tokens = tokenize_query(query)
|
|
boosted_sources = set(keyword_boost(query))
|
|
scored: list[tuple[float, int]] = []
|
|
|
|
for i, doc in enumerate(metadata):
|
|
title = (doc.get("title") or "").lower()
|
|
source = (doc.get("source") or "").lower()
|
|
text = (doc.get("text") or "").lower()
|
|
hay = f"{title}\n{source}\n{text}"
|
|
|
|
score = 0.0
|
|
|
|
# 领域词表强加权
|
|
if doc.get("source") in boosted_sources:
|
|
score += 20.0
|
|
|
|
# 完整问题/短语命中
|
|
if query_lower and query_lower in hay:
|
|
score += 6.0
|
|
if query_lower and query_lower in title:
|
|
score += 5.0
|
|
|
|
# token 频次 + 字段权重
|
|
for token in tokens:
|
|
if token in title:
|
|
score += 4.0
|
|
if token in source:
|
|
score += 3.0
|
|
count = text.count(token)
|
|
if count:
|
|
# 类 BM25：频次递减，避免长文刷分
|
|
score += min(3.0, 1.0 + count ** 0.5)
|
|
|
|
# 配置/安装类问题：代码块和关键配置变量优先，避免模型拿不到原文示例后自行编造
|
|
config_query = any(k in query_lower for k in ["compose", "docker", "postgres", "postgresql", "安装", "部署", "配置", "环境变量"])
|
|
if config_query and "[代码块]" in text:
|
|
score += 10.0
|
|
config_hits = sum(1 for k in [
|
|
"image:", "services:", "environment:", "volumes:", "ports:",
|
|
"postgres_db", "postgres_user", "postgres_password",
|
|
"email", "token", "go_web_port", "database", "db_",
|
|
] if k in text)
|
|
score += min(20.0, config_hits * 3.0)
|
|
|
|
# 查询中英文/数字 token 对安装、docker、postgresql 等问题很关键
|
|
exact_tokens = re.findall(r"[a-z0-9][a-z0-9_.:-]{1,}", query_lower)
|
|
for token in exact_tokens:
|
|
if token in hay:
|
|
score += 2.5
|
|
|
|
if score > 0:
|
|
scored.append((score, i))
|
|
|
|
scored.sort(reverse=True, key=lambda x: x[0])
|
|
|
|
docs = []
|
|
for score, idx in scored:
|
|
doc = metadata[idx].copy()
|
|
doc["score"] = float(score)
|
|
src_count = sum(1 for d in docs if d["source"] == doc["source"])
|
|
if src_count < 2:
|
|
docs.append(doc)
|
|
if len(docs) >= top_k:
|
|
break
|
|
|
|
return docs
|
|
|
|
|
|
def vector_search(query: str, top_k: int = TOP_K) -> list[dict]:
|
|
"""向量检索：仅开发/宿主机模式使用，需要 sentence-transformers。"""
|
|
embeddings, metadata = load_index()
|
|
if len(embeddings) == 0:
|
|
return []
|
|
|
|
query_vec = get_model().encode([query], normalize_embeddings=True)
|
|
scores = (embeddings @ query_vec.T).flatten()
|
|
|
|
boosted_sources = keyword_boost(query)
|
|
if boosted_sources:
|
|
for i, doc in enumerate(metadata):
|
|
if doc["source"] in boosted_sources:
|
|
scores[i] += 0.3
|
|
|
|
query_words = set(re.findall(r'[\u4e00-\u9fff]+', query))
|
|
for i, doc in enumerate(metadata):
|
|
text_lower = doc["text"].lower()
|
|
matches = sum(1 for w in query_words if w in text_lower)
|
|
if matches >= 2:
|
|
scores[i] += 0.15 * matches
|
|
|
|
top_indices = np.argsort(scores)[::-1][:top_k]
|
|
|
|
docs = []
|
|
for idx in top_indices:
|
|
doc = metadata[idx].copy()
|
|
doc["score"] = float(scores[idx])
|
|
src_count = sum(1 for d in docs if d["source"] == doc["source"])
|
|
if src_count < 2:
|
|
docs.append(doc)
|
|
return docs
|
|
|
|
|
|
def expand_context(docs: list[dict], window: int = 1, max_chars: int = 12000, max_sources: int = 5) -> list[dict]:
|
|
"""为命中的 chunk 补充同源相邻上下文，减少回答过短/断章取义。"""
|
|
if not docs:
|
|
return []
|
|
|
|
metadata = load_metadata()
|
|
by_id = {d.get("id"): d for d in metadata}
|
|
expanded: list[dict] = []
|
|
seen: set[str] = set()
|
|
total_chars = 0
|
|
|
|
def parse_id(doc_id: str) -> tuple[str, int] | None:
|
|
if not doc_id or "#" not in doc_id:
|
|
return None
|
|
src, idx = doc_id.rsplit("#", 1)
|
|
if not idx.isdigit():
|
|
return None
|
|
return src, int(idx)
|
|
|
|
allowed_sources = list(dict.fromkeys(d.get("source") for d in docs if d.get("source")))[:max_sources]
|
|
allowed_sources_set = set(allowed_sources)
|
|
|
|
for doc in docs:
|
|
if doc.get("source") not in allowed_sources_set:
|
|
continue
|
|
parsed = parse_id(doc.get("id", ""))
|
|
candidates = [doc]
|
|
if parsed:
|
|
src, idx = parsed
|
|
neighbors = []
|
|
for offset in range(-window, window + 1):
|
|
item = by_id.get(f"{src}#{idx + offset}")
|
|
if item:
|
|
neighbors.append(item)
|
|
candidates = neighbors or [doc]
|
|
|
|
for item in candidates:
|
|
doc_id = item.get("id", "")
|
|
if doc_id in seen:
|
|
continue
|
|
text_len = len(item.get("text", ""))
|
|
if total_chars + text_len > max_chars and expanded:
|
|
continue
|
|
copy = item.copy()
|
|
copy.setdefault("score", doc.get("score", 0))
|
|
expanded.append(copy)
|
|
seen.add(doc_id)
|
|
total_chars += text_len
|
|
|
|
return expanded
|
|
|
|
|
|
def search(query: str, top_k: int = TOP_K, expand: bool = True, context_window: int = 1) -> list[dict]:
|
|
mode = os.environ.get("KNOWLEDGE_SEARCH_MODE", "lite").lower()
|
|
if mode in {"vector", "embedding", "semantic"}:
|
|
docs = vector_search(query, top_k=top_k)
|
|
else:
|
|
docs = lite_search(query, top_k=top_k)
|
|
return expand_context(docs, window=context_window) if expand else docs
|
|
|
|
|
|
def get_stats() -> dict:
|
|
metadata = load_metadata()
|
|
cache = load_qa_cache()
|
|
return {
|
|
"total_chunks": len(metadata),
|
|
"cached_answers": len(cache),
|
|
"site_url": SITE_URL,
|
|
"search_mode": os.environ.get("KNOWLEDGE_SEARCH_MODE", "lite"),
|
|
"embedding_model": EMBED_MODEL,
|
|
"cache_version": CACHE_VERSION,
|
|
}
|
|
|
|
|
|
# ============ CLI ============
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
if len(sys.argv) > 1 and sys.argv[1] == "rebuild":
|
|
rebuild_index()
|
|
elif len(sys.argv) > 1 and sys.argv[1] == "search":
|
|
query = " ".join(sys.argv[2:])
|
|
results = search(query)
|
|
for r in results:
|
|
print(f"\n📄 {r['title']} ({r['source']}) [分数:{r.get('score', 0):.3f}]")
|
|
print(f" URL: {r['url']}")
|
|
print(f" {r['text'][:200]}...")
|
|
elif len(sys.argv) > 1 and sys.argv[1] == "stats":
|
|
print(json.dumps(get_stats(), ensure_ascii=False, indent=2))
|
|
elif len(sys.argv) > 1 and sys.argv[1] == "cache":
|
|
cache = load_qa_cache()
|
|
for k, v in cache.items():
|
|
print(f"\n🔑 {k[:8]}... Q: {v['question'][:50]}")
|
|
else:
|
|
print("用法:")
|
|
print(" python -m knowledge rebuild - 重建索引")
|
|
print(" python -m knowledge search <查询> - 搜索")
|
|
print(" python -m knowledge stats - 统计信息")
|
|
print(" python -m knowledge cache - 查看缓存")
|