feat: rebuild tg spam guard v2 architecture
This commit is contained in:
ngfchl
@@ -0,0 +1,719 @@
|
||||
"""知识库管理模块 - 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 - 查看缓存")
|
||||
Reference in New Issue
Block a user