T3: TreeSearch 结构化文档两级检索 (build_kb_tree + tree_search BM25+向量融合, 接入 kb_router 为 tree_chunk 源, 新增 GET /api/knowledge/tree-search)
This commit is contained in:
QihuAI Design committed
1 parent
e30e385f11
commit
be7d792dc2
3 files changed
+2124
-1622
No files matched your search
@@ -0,0 +1,405 @@
|
||||
"""
|
||||
TreeSearch — 结构化文档两级检索
|
||||
================================
|
||||
|
||||
为知识库提供有结构的「两级检索」能力（零额外基础设施成本，复用 MySQL 已存的分段数据）：
|
||||
|
||||
Level 1 (build_kb_tree) : 知识库 -> 目录 -> 文档 的嵌套结构树，用于浏览/定位/面包屑。
|
||||
Level 2 (tree_search) : 在给定作用域（整个 KB / 某目录 / 某文档）内做分段(chunk)级检索，
|
||||
BM25 关键词排序 + 可选 pgvector 语义向量 boost，返回带完整
|
||||
面包屑(KB名 -> 目录路径 -> 文件名 -> 段序号/标题)的精准命中。
|
||||
|
||||
设计要点：
|
||||
* Level 2 直接消费 knowledge_document_chunks 表（content + chunk_metadata.title/keywords），
|
||||
因此召回的是「文档内部的精确片段」，而非整篇文档——这正是 RAG 想要的细粒度上下文。
|
||||
* 与 kb_router_search 融合：tree_search 作为 RRF 的一个新来源 (source="tree_chunk") 接入，
|
||||
聊天 KB 上下文自动获得结构化分片段命中。
|
||||
* 向量 boost 为可选且优雅降级：当 pgvector 无对应数据时，自动退化为纯 BM25，不影响主链路。
|
||||
"""
|
||||
|
||||
import logging
|
||||
import math
|
||||
from typing import List, Dict, Optional
|
||||
|
||||
logger = logging.getLogger("aiflow.tree_search")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 中文友好的分词（jieba 优先，否则按字符 bigram + 英文正则兜底）
|
||||
# ---------------------------------------------------------------------------
|
||||
def _tokenize(text: str) -> List[str]:
|
||||
if not text:
|
||||
return []
|
||||
toks: List[str] = []
|
||||
try:
|
||||
import re
|
||||
# 英文/数字 token
|
||||
for m in re.findall(r"[a-zA-Z0-9_]+", text.lower()):
|
||||
if len(m) >= 1:
|
||||
toks.append(m)
|
||||
# 中文按 jieba 或字符切分
|
||||
cjk = "".join(re.findall(r"[一-鿿]", text))
|
||||
if cjk:
|
||||
try:
|
||||
import jieba
|
||||
for w in jieba.lcut(cjk):
|
||||
w = w.strip().lower()
|
||||
if w:
|
||||
toks.append(w)
|
||||
except Exception:
|
||||
# 字符级 bigram（兼顾召回与精度）
|
||||
for i in range(len(cjk)):
|
||||
toks.append(cjk[i])
|
||||
if i + 1 < len(cjk):
|
||||
toks.append(cjk[i:i + 2])
|
||||
except Exception:
|
||||
toks = [t for t in text.lower().split() if t]
|
||||
return toks
|
||||
|
||||
|
||||
def _chunk_display_text(chunk) -> str:
|
||||
"""拼装分段的检索文本：正文 + 标题 + 关键词（来自 chunk_metadata）。"""
|
||||
parts = []
|
||||
content = getattr(chunk, "content", None) or ""
|
||||
parts.append(content)
|
||||
md = getattr(chunk, "chunk_metadata", None) or {}
|
||||
if isinstance(md, dict):
|
||||
title = md.get("title") or md.get("heading") or ""
|
||||
keywords = md.get("keywords") or []
|
||||
if title:
|
||||
parts.append(title)
|
||||
if isinstance(keywords, list) and keywords:
|
||||
parts.append(" ".join(str(k) for k in keywords))
|
||||
elif isinstance(keywords, str) and keywords:
|
||||
parts.append(keywords)
|
||||
return "\n".join(p for p in parts if p).strip()
|
||||
|
||||
|
||||
def _chunk_title(chunk) -> str:
|
||||
md = getattr(chunk, "chunk_metadata", None) or {}
|
||||
if isinstance(md, dict):
|
||||
return md.get("title") or md.get("heading") or ""
|
||||
return ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# BM25（在作用域内的分段语料上即时计算 IDF）
|
||||
# ---------------------------------------------------------------------------
|
||||
def _bm25_rank(query_tokens: List[str], corpus: List[Dict], k1: float = 1.5, b: float = 0.75) -> List[Dict]:
|
||||
"""corpus: [{"id":..., "tokens":[...], "chunk":<orm>, "text":str}] -> 返回带 bm25 分的列表"""
|
||||
if not query_tokens or not corpus:
|
||||
return []
|
||||
n = len(corpus)
|
||||
df = {}
|
||||
for qt in set(query_tokens):
|
||||
df[qt] = sum(1 for d in corpus if qt in d["tokens"])
|
||||
idf = {qt: math.log((n - df[qt] + 0.5) / (df[qt] + 0.5) + 1) for qt in query_tokens}
|
||||
avgdl = sum(len(d["tokens"]) for d in corpus) / max(n, 1)
|
||||
scored = []
|
||||
for d in corpus:
|
||||
dl = len(d["tokens"])
|
||||
score = 0.0
|
||||
for qt in set(query_tokens):
|
||||
if qt in d["tokens"]:
|
||||
f = d["tokens"].count(qt)
|
||||
score += idf[qt] * (f * (k1 + 1)) / (f + k1 * (1 - b + b * (dl / max(avgdl, 1))))
|
||||
if score > 0:
|
||||
scored.append({**d, "bm25": score})
|
||||
scored.sort(key=lambda x: x["bm25"], reverse=True)
|
||||
return scored
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Level 1：知识库目录树
|
||||
# ---------------------------------------------------------------------------
|
||||
def build_kb_tree(db, user_id: Optional[str] = None, kb_id: Optional[str] = None) -> List[Dict]:
|
||||
"""构造 知识库 -> 目录(可嵌套) -> 文档 的结构树。"""
|
||||
from app.models import KnowledgeBase, KnowledgeDirectory, KnowledgeDocument
|
||||
|
||||
q = db.query(KnowledgeBase)
|
||||
if user_id:
|
||||
q = q.filter(KnowledgeBase.user_id == user_id)
|
||||
if kb_id:
|
||||
q = q.filter(KnowledgeBase.id == kb_id)
|
||||
bases = q.all()
|
||||
|
||||
tree = []
|
||||
for kb in bases:
|
||||
dirs = db.query(KnowledgeDirectory).filter(
|
||||
KnowledgeDirectory.knowledge_base_id == kb.id
|
||||
).all()
|
||||
docs = db.query(KnowledgeDocument).filter(
|
||||
KnowledgeDocument.knowledge_base_id == kb.id,
|
||||
KnowledgeDocument.status == "completed",
|
||||
).all()
|
||||
dir_by_id = {d.id: d for d in dirs}
|
||||
node_by_id = {}
|
||||
|
||||
def make_dir_node(d):
|
||||
return {
|
||||
"id": d.id,
|
||||
"name": d.name,
|
||||
"type": "directory",
|
||||
"children": [],
|
||||
"documents": [],
|
||||
}
|
||||
|
||||
for d in dirs:
|
||||
node_by_id[d.id] = make_dir_node(d)
|
||||
root_dirs = []
|
||||
for d in dirs:
|
||||
node = node_by_id[d.id]
|
||||
parent = d.parent_id
|
||||
if parent and parent in node_by_id:
|
||||
node_by_id[parent]["children"].append(node)
|
||||
else:
|
||||
root_dirs.append(node)
|
||||
for doc in docs:
|
||||
item = {
|
||||
"id": doc.id,
|
||||
"name": doc.filename,
|
||||
"type": "document",
|
||||
"file_type": doc.file_type,
|
||||
"chunk_count": doc.chunk_count or 0,
|
||||
}
|
||||
if doc.directory_id and doc.directory_id in node_by_id:
|
||||
node_by_id[doc.directory_id]["documents"].append(item)
|
||||
else:
|
||||
root_dirs.append(item) # 顶层文档
|
||||
tree.append({
|
||||
"id": kb.id,
|
||||
"name": kb.name,
|
||||
"type": "knowledge_base",
|
||||
"directories": root_dirs,
|
||||
"document_count": len(docs),
|
||||
})
|
||||
return tree
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 面包屑：KB名 -> 目录路径 -> 文件名 -> 段标题/序号
|
||||
# ---------------------------------------------------------------------------
|
||||
def _breadcrumb(db, kb_id: str, directory_id: Optional[str], doc_id: str, chunk_index: int, title: str) -> Dict:
|
||||
from app.models import KnowledgeBase, KnowledgeDirectory, KnowledgeDocument
|
||||
kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id).first()
|
||||
kb_name = kb.name if kb else kb_id
|
||||
path = []
|
||||
cur = directory_id
|
||||
guard = 0
|
||||
while cur and guard < 20:
|
||||
d = db.query(KnowledgeDirectory).filter(KnowledgeDirectory.id == cur).first()
|
||||
if not d:
|
||||
break
|
||||
path.insert(0, d.name)
|
||||
cur = d.parent_id
|
||||
guard += 1
|
||||
doc = db.query(KnowledgeDocument).filter(KnowledgeDocument.id == doc_id).first()
|
||||
doc_name = doc.filename if doc else doc_id
|
||||
return {
|
||||
"knowledge_base": kb_name,
|
||||
"directory_path": "/".join(path) if path else "/",
|
||||
"document": doc_name,
|
||||
"chunk_index": chunk_index,
|
||||
"chunk_title": title,
|
||||
"label": f"{kb_name}/{('/'.join(path) if path else '')}/{doc_name}" + (f"#段{chunk_index}" + (f"·{title}" if title else "")),
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Level 2：结构化分片段检索
|
||||
# ---------------------------------------------------------------------------
|
||||
async def tree_search(
|
||||
db,
|
||||
query: str,
|
||||
user_id: Optional[str] = None,
|
||||
kb_ids: Optional[List[str]] = None,
|
||||
scope_base: Optional[str] = None,
|
||||
scope_dir: Optional[str] = None,
|
||||
scope_doc: Optional[str] = None,
|
||||
top_k: int = 5,
|
||||
mode: str = "hybrid",
|
||||
max_tokens: int = 1200,
|
||||
) -> List[Dict]:
|
||||
"""
|
||||
两级检索的 Level 2：在作用域内做分片段精准检索。
|
||||
|
||||
Args:
|
||||
query : 查询文本
|
||||
kb_ids : 限定知识库范围（多个）
|
||||
scope_base : 限定单一知识库
|
||||
scope_dir : 限定单一目录（含子目录）
|
||||
scope_doc : 限定单一文档
|
||||
top_k : 返回命中数
|
||||
mode : "bm25" | "vector" | "hybrid"(默认)
|
||||
Returns:
|
||||
[{
|
||||
"content": <分段文本>,
|
||||
"score": 0-1,
|
||||
"source": "tree_chunk",
|
||||
"metadata": { "breadcrumb": {...}, "snippet": ..., "chunk_index": int, "title": str }
|
||||
}, ...]
|
||||
"""
|
||||
from app.models import KnowledgeBase, KnowledgeDirectory, KnowledgeDocument, KnowledgeDocumentChunk
|
||||
|
||||
if not query or not query.strip():
|
||||
return []
|
||||
|
||||
# 1) 解析作用域 -> 目标文档 id 集合
|
||||
target_kb_ids = list(kb_ids or [])
|
||||
if scope_base and scope_base not in target_kb_ids:
|
||||
target_kb_ids.append(scope_base)
|
||||
|
||||
doc_scope = set()
|
||||
if scope_doc:
|
||||
doc_scope.add(scope_doc)
|
||||
if scope_dir:
|
||||
# 收集该目录及其所有子目录下的文档
|
||||
dirs = db.query(KnowledgeDirectory).all()
|
||||
dir_by_id = {d.id: d for d in dirs}
|
||||
sub = set()
|
||||
|
||||
def collect(did):
|
||||
sub.add(did)
|
||||
for d in dirs:
|
||||
if d.parent_id == did:
|
||||
collect(d.id)
|
||||
collect(scope_dir)
|
||||
docs = db.query(KnowledgeDocument).filter(KnowledgeDocument.directory_id.in_(sub)).all()
|
||||
for d in docs:
|
||||
doc_scope.add(d.id)
|
||||
if not doc_scope and not target_kb_ids and not user_id:
|
||||
return []
|
||||
|
||||
# 2) 取分段语料
|
||||
q = db.query(KnowledgeDocumentChunk)
|
||||
if doc_scope:
|
||||
q = q.filter(KnowledgeDocumentChunk.document_id.in_(doc_scope))
|
||||
else:
|
||||
if target_kb_ids:
|
||||
q = q.filter(KnowledgeDocumentChunk.knowledge_base_id.in_(target_kb_ids))
|
||||
elif user_id:
|
||||
kb_ids_u = [r[0] for r in db.query(KnowledgeBase.id).filter(KnowledgeBase.user_id == user_id).all()]
|
||||
if not kb_ids_u:
|
||||
return []
|
||||
q = q.filter(KnowledgeDocumentChunk.knowledge_base_id.in_(kb_ids_u))
|
||||
chunks = q.order_by(KnowledgeDocumentChunk.chunk_index).limit(500).all()
|
||||
if not chunks:
|
||||
return []
|
||||
|
||||
query_tokens = _tokenize(query)
|
||||
corpus = []
|
||||
for c in chunks:
|
||||
text = _chunk_display_text(c)
|
||||
if len(text.strip()) < 8:
|
||||
continue
|
||||
corpus.append({
|
||||
"id": c.id,
|
||||
"tokens": _tokenize(text),
|
||||
"chunk": c,
|
||||
"text": text,
|
||||
})
|
||||
|
||||
# 3) BM25 排序
|
||||
bm25_hits = _bm25_rank(query_tokens, corpus) if mode in ("bm25", "hybrid") else []
|
||||
|
||||
# 4) 可选向量 boost
|
||||
vector_boost = {}
|
||||
if mode in ("vector", "hybrid"):
|
||||
try:
|
||||
from app.pgvector_search import vector_search
|
||||
for kid in (target_kb_ids or [None]):
|
||||
try:
|
||||
vres = await vector_search(query, kid, top_k=top_k * 3)
|
||||
except Exception as e:
|
||||
logger.warning(f"[tree_search] vector_search failed: {e}")
|
||||
continue
|
||||
for vr in vres:
|
||||
vcontent = (vr.get("content") or "").strip()
|
||||
if vcontent:
|
||||
vector_boost[vcontent] = max(vector_boost.get(vcontent, 0.0), float(vr.get("score", 0.0)))
|
||||
except Exception as e:
|
||||
logger.debug(f"[tree_search] vector unavailable: {e}")
|
||||
|
||||
# 5) 融合
|
||||
merged_map = {}
|
||||
for rank, h in enumerate(bm25_hits):
|
||||
key = h["text"][:300]
|
||||
boost = 0.0
|
||||
for vc, vs in vector_boost.items():
|
||||
if vc and vc in h["text"]:
|
||||
boost = max(boost, vs)
|
||||
merged_map[key] = {
|
||||
"chunk": h["chunk"],
|
||||
"bm25": h["bm25"],
|
||||
"vec": boost,
|
||||
"rank": rank + 1,
|
||||
}
|
||||
# 纯向量模式：把仅在向量里命中的也纳入
|
||||
if mode == "vector":
|
||||
for vc, vs in vector_boost.items():
|
||||
if vc[:300] not in merged_map:
|
||||
# 找到对应 chunk
|
||||
for c in corpus:
|
||||
if vc in c["text"]:
|
||||
merged_map[vc[:300]] = {"chunk": c["chunk"], "bm25": 0.0, "vec": vs, "rank": 999}
|
||||
break
|
||||
|
||||
if not merged_map:
|
||||
return []
|
||||
|
||||
# 综合分数（BM25 主导，向量做 0~1 归一 boost）
|
||||
max_bm25 = max((m["bm25"] for m in merged_map.values()), default=1.0) or 1.0
|
||||
items = []
|
||||
for key, m in merged_map.items():
|
||||
c = m["chunk"]
|
||||
norm_bm25 = (m["bm25"] / max_bm25) if max_bm25 > 0 else 0.0
|
||||
norm_vec = min(m["vec"], 1.0)
|
||||
combined = norm_bm25 * 0.7 + norm_vec * 0.3
|
||||
title = _chunk_title(c)
|
||||
breadcrumb = _breadcrumb(db, c.knowledge_base_id, c.document_id and _doc_dir(db, c.document_id), c.document_id, c.chunk_index, title)
|
||||
snippet = _make_snippet(c.content or "", query)
|
||||
items.append({
|
||||
"content": c.content or "",
|
||||
"score": round(combined, 4),
|
||||
"source": "tree_chunk",
|
||||
"metadata": {
|
||||
"breadcrumb": breadcrumb,
|
||||
"snippet": snippet,
|
||||
"chunk_index": c.chunk_index,
|
||||
"title": title,
|
||||
"document_id": c.document_id,
|
||||
"knowledge_base_id": c.knowledge_base_id,
|
||||
},
|
||||
})
|
||||
items.sort(key=lambda x: x["score"], reverse=True)
|
||||
return items[:top_k]
|
||||
|
||||
|
||||
def _doc_dir(db, doc_id: str):
|
||||
from app.models import KnowledgeDocument
|
||||
d = db.query(KnowledgeDocument).filter(KnowledgeDocument.id == doc_id).first()
|
||||
return d.directory_id if d else None
|
||||
|
||||
|
||||
def _make_snippet(content: str, query: str, win: int = 80) -> str:
|
||||
if not content:
|
||||
return ""
|
||||
q = query.strip()
|
||||
idx = content.find(q)
|
||||
if idx < 0:
|
||||
# 尝试首个 token
|
||||
toks = _tokenize(q)
|
||||
for t in toks:
|
||||
idx = content.find(t)
|
||||
if idx >= 0:
|
||||
break
|
||||
if idx < 0:
|
||||
return content[:win * 2]
|
||||
start = max(0, idx - win)
|
||||
end = min(len(content), idx + len(q) + win)
|
||||
snip = content[start:end]
|
||||
if start > 0:
|
||||
snip = "..." + snip
|
||||
if end < len(content):
|
||||
snip = snip + "..."
|
||||
return snip
|
||||
Reference in new issue
Block a user