406 lines
15 KiB
Python
406 lines
15 KiB
Python
"""
|
|
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
|