T3: TreeSearch 结构化文档两级检索 (build_kb_tree + tree_search BM25+向量融合, 接入 kb_router 为 tree_chunk 源, 新增 GET /api/knowledge/tree-search)

backend/python/app/api/knowledge.py

@@ -1,3 +1,4 @@
+import asyncio
 from fastapi import APIRouter, Depends, HTTPException, Query, UploadFile, File, Form
 from sqlalchemy.orm import Session
 from sqlalchemy import and_, desc, asc
@@ -851,6 +852,15 @@ def get_documents(
 
 
 @router.post("/bases/{base_id}/documents")
+async def _index_kb_doc_to_pgvector(doc_id: str, kb_id: str, content: str):
+    """将文档内容异步索引到 pgvector，供语义向量检索使用（best-effort）。"""
+    try:
+        from app.pgvector_search import index_document
+        await index_document(doc_id, kb_id, content, metadata={"source": "kb_upload"})
+    except Exception as e:
+        logger.warning(f"[kb] 向量索引失败 doc={doc_id}: {e}")
+
+
 async def upload_document(
     base_id: str,
     file: Optional[UploadFile] = File(None),
@@ -961,6 +971,12 @@ async def upload_document(
     db.commit()
     db.refresh(doc)
 
+    # T1: 异步索引到 pgvector 语义库（best-effort，不阻塞上传响应）
+    try:
+        asyncio.create_task(_index_kb_doc_to_pgvector(doc.id, base_id, content))
+    except Exception:
+        pass
+
     LogService.log_operation(
         user_id=current_user.id, username=current_user.username, action="upload",
         target="knowledge_document", target_id=doc.id,
@@ -1415,6 +1431,58 @@ def search_knowledge(
     }
 
 
+@router.get("/tree-search")
+async def tree_search_endpoint(
+    query: str = Query(min_length=1, description="检索查询"),
+    kb_id: Optional[str] = Query(default=None, description="限定知识库ID"),
+    scope: Optional[str] = Query(
+        default=None,
+        description="作用域: doc:<doc_id> | dir:<dir_id> | base:<base_id>",
+    ),
+    top_k: int = Query(default=5, ge=1, le=20),
+    db: Session = Depends(get_db),
+    current_user: User = Depends(get_current_user),
+):
+    """TreeSearch 结构化文档两级检索。
+
+    Level 1: 返回用户可见的知识库 -> 目录 -> 文档结构树(tree)。
+    Level 2: 在作用域内做分片段(chunk)级精准检索，命中带完整面包屑与片段摘要。
+    """
+    from app.tree_search import build_kb_tree, tree_search
+
+    # Level 1: 结构树
+    tree = build_kb_tree(db, user_id=current_user.id, kb_id=kb_id)
+
+    # Level 2: 分片段检索
+    scope_doc = scope_dir = scope_base = None
+    if scope:
+        if scope.startswith("doc:"):
+            scope_doc = scope[len("doc:"):]
+        elif scope.startswith("dir:"):
+            scope_dir = scope[len("dir:"):]
+        elif scope.startswith("base:"):
+            scope_base = scope[len("base:"):]
+
+    hits = await tree_search(
+        db, query,
+        user_id=current_user.id,
+        kb_ids=[kb_id] if kb_id else None,
+        scope_base=scope_base, scope_dir=scope_dir, scope_doc=scope_doc,
+        top_k=top_k, mode="hybrid",
+    )
+
+    return {
+        "code": 0,
+        "data": {
+            "query": query,
+            "tree": tree,
+            "hits": hits,
+            "total": len(hits),
+        },
+        "message": "ok",
+    }
+
+
 def generate_document_excerpt(doc, query: str) -> str:
     """根据文档类型生成模拟的原文摘录"""
     excerpts = {

backend/python/app/kb_router.py

@@ -67,11 +67,13 @@ async def _search_external(service: Dict, query: str, top_k: int = 5) -> List[Di
 async def kb_router_search(
     query: str,
     kb_id: Optional[str] = None,
+    kb_ids: Optional[List[str]] = None,
     top_k: int = 5,
     bm25_results: Optional[List[Dict]] = None,
     db_session=None,
     use_vector: bool = True,
     use_external: bool = True,
+    use_tree: bool = True,
 ) -> List[Dict]:
     """
     知识库路由层 — 混合检索
@@ -79,7 +81,8 @@ async def kb_router_search(
     检索策略:
     1. 本地 BM25 检索 (MySQL, 已有)
     2. 向量语义检索 (PostgreSQL+pgvector, 新增)
-    3. 外部KB检索 (Dify, 按配置)
+    3. 结构化分片段检索 (TreeSearch, 新增: 知识库目录树 + chunk 级精准命中)
+    4. 外部KB检索 (Dify, 按配置)
 
     融合方式: Reciprocal Rank Fusion (RRF)
 
@@ -88,9 +91,10 @@ async def kb_router_search(
         kb_id: 知识库ID
         top_k: 返回结果数
         bm25_results: BM25检索结果
-        db_session: 数据库会话(用于获取外部KB配置)
+        db_session: 数据库会话(用于获取外部KB配置与分片段检索)
         use_vector: 是否使用向量检索
         use_external: 是否使用外部KB检索
+        use_tree: 是否使用 TreeSearch 结构化分片段检索
 
     Returns:
         融合后的检索结果列表
@@ -112,7 +116,13 @@ async def kb_router_search(
     if use_vector:
         try:
             from app.pgvector_search import vector_search
-            vec_results = await vector_search(query, kb_id, top_k)
+            kb_targets = kb_ids if kb_ids else ([kb_id] if kb_id else [None])
+            vec_results = []
+            for _kid in kb_targets:
+                try:
+                    vec_results.extend(await vector_search(query, _kid, top_k))
+                except Exception as e:
+                    logger.warning(f"Vector search failed for kb={_kid}: {e}")
             for rank, item in enumerate(vec_results):
                 all_results.append({
                     "content": item.get("content", ""),
@@ -126,6 +136,25 @@ async def kb_router_search(
         except Exception as e:
             logger.warning(f"Vector search failed: {e}")
 
+    # 2.5 结构化分片段检索 (TreeSearch, Level 2)
+    if use_tree and db_session is not None:
+        try:
+            from app.tree_search import tree_search
+            kb_targets = kb_ids if kb_ids else ([kb_id] if kb_id else None)
+            tree_hits = await tree_search(
+                db_session, query, kb_ids=kb_targets, top_k=top_k, mode="hybrid",
+            )
+            for rank, item in enumerate(tree_hits):
+                all_results.append({
+                    "content": item.get("content", ""),
+                    "score": item.get("score", 0.5),
+                    "source": "tree_chunk",
+                    "rank": rank + 1,
+                    "metadata": item.get("metadata", {}),
+                })
+        except Exception as e:
+            logger.warning(f"TreeSearch failed: {e}")
+
     # 3. 外部KB检索 (Dify)
     if use_external and db_session:
         services = await _get_external_services(db_session)

backend/python/app/tree_search.py

@@ -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