From 913380f738bfe7d92baaae48b15b9c059ec8f790 Mon Sep 17 00:00:00 2001 From: ThearchyHelios Date: Sat, 19 Sep 2026 09:44:09 +0800 Subject: [PATCH] =?UTF-8?q?=E5=80=92=E6=8E=92=E7=B4=A2=E5=BC=95=E5=8F=96?= =?UTF-8?q?=E4=BB=A3=E5=85=A8=E8=A1=A8=E9=87=8D=E7=AE=97=EF=BC=9BPDF/Word?= =?UTF-8?q?=20=E7=BB=88=E4=BA=8E=E4=BC=A0=E5=BE=97=E8=BF=9B=E6=9D=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit **倒排**：以前每次检索都 select(Chunk) 全量载入，再在 Python 里重建一遍 BM25 ——不只是扫表，是对每一段重新分词、重新统计词频，每次查询都来一遍。几百段 无感，上万段就是秒级，而那正是知识库开始有用的规模。 没用 SQLite 的 FTS5。方案里写了实施前先验证，验了：trigram 分词器要求查询 至少三个字符，中文的「商家」「订单」「权限」这些两字词一个都召不回。照搬会 让召回大幅倒退。现有分词器本来就做中文单字 + 二元组，自己存一张 chunk_terms 既保住召回，也不把知识库锁死在 SQLite 上。 快是次要的，一致才是前提——所以先写了"倒排必须和全表扫给出同一个名次"的 测试。它当场抓到两个真问题： 1. BM25 被算了两次：倒排按全集统计（N、avg_len、df）算一次，hybrid_rank 又在候选子集上重建一次，而子集的 df 和平均长度是另一套数。改成把倒排 算好的分传下去，不再重算。 2. 原来的 BM25.score 遍历的是**带重复**的查询词表：「导出订单最多多少行」 里「多」出现两次，贡献就加两次。倒排用 set() 去重，顶部那条差了 2.32。 改成按查询词频加权。 还有一处是排序确定性：分数并列时稳定排序保留输入顺序，而两条路径喂进去的 顺序天生不同——同一个查询会因为走了哪条路给出不同名次。两边都按 id 定序。 倒排跟着 ingest 建、跟着删文档清、跟着 reindex 重建；没建倒排的片段（存量 数据升级上来就是这样）退回全表扫——宁可慢，也不要因为索引没建好就说"没搜到"。 **解析**：upload 以前只收 UTF-8，PDF 传上去直接 415，而 mime 字段记了却从不 参与判断。现在认 PDF / DOCX / HTML。解析库走可选依赖 [docs]（几十兆，只用 纯文本的人不该为此付出安装时间），HTML 用已在依赖里的 bs4。 几处是实际验出来的，不是想出来的：DOCX 要连表格一起取（合同和口径说明的关键 信息常常整个在表里）；PDF 每行拖着排版填充空格，不清掉会混进切块和词频统计 （实测 416 字里只有 68 字是正文）；扫描件提不出文字要明说，否则悄悄存一份 空文档，用户以为传成功了却永远搜不到。 PDF/DOCX 两条路径在临时 venv 里用真文件验过（含中文、表格、扫描件、坏文件）， 没有动开发环境。 --- backend/app/api/knowledge.py | 14 ++- backend/app/db/base.py | 1 + backend/app/db/models.py | 26 +++++ backend/app/memory/embeddings.py | 9 +- backend/app/memory/inverted.py | 131 ++++++++++++++++++++++++ backend/app/memory/kb.py | 93 +++++++++++++---- backend/app/memory/parsing.py | 120 ++++++++++++++++++++++ backend/pyproject.toml | 3 + backend/tests/test_kb.py | 75 +++++++++++++- backend/tests/test_retrieval_quality.py | 90 ++++++++++++++++ 10 files changed, 535 insertions(+), 27 deletions(-) create mode 100644 backend/app/memory/inverted.py create mode 100644 backend/app/memory/parsing.py diff --git a/backend/app/api/knowledge.py b/backend/app/api/knowledge.py index 60802dd2..37c87255 100644 --- a/backend/app/api/knowledge.py +++ b/backend/app/api/knowledge.py @@ -9,7 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from app.db.base import get_session from app.db.models import Document, MemoryItem, Skill -from app.memory import kb, store +from app.memory import inverted, kb, store # -------------------------------------------------------------------------- # 长期记忆 @@ -153,10 +153,14 @@ async def upload( raw = await file.read() if len(raw) > 10 * 1024 * 1024: raise HTTPException(413, "文件超过 10MB") + + from app.memory.parsing import UnsupportedDocument, extract + try: - content = raw.decode("utf-8") - except UnicodeDecodeError: - raise HTTPException(415, "目前只支持 UTF-8 文本文件（txt / md / json / csv 等）") from None + content = extract(raw, file.filename or "", file.content_type or "") + except UnsupportedDocument as e: + # 原样交回：里面写的是"装哪个包"或"这是扫描件"，都能照着做 + raise HTTPException(415, str(e)) from None return await kb.ingest_document( session, collection=collection, @@ -205,6 +209,8 @@ async def embedding_status( "stale_chunks": await kb.stale_count(session, collection), "has_semantics": has_semantics(), "default_alpha": default_alpha(), + # 没建倒排的片段会走全表扫——结果对，但慢 + "unindexed_chunks": await inverted.missing_count(session, collection), } diff --git a/backend/app/db/base.py b/backend/app/db/base.py index 43f9a7db..83e5d130 100644 --- a/backend/app/db/base.py +++ b/backend/app/db/base.py @@ -81,6 +81,7 @@ _COLUMN_MIGRATIONS: list[tuple[str, str, str]] = [ ("chunks", "embed_dim", "INTEGER DEFAULT 0"), ("memories", "embed_model", "VARCHAR(100) DEFAULT ''"), ("memories", "embed_dim", "INTEGER DEFAULT 0"), + ("chunks", "token_len", "INTEGER DEFAULT 0"), ] diff --git a/backend/app/db/models.py b/backend/app/db/models.py index 70c39b5e..7529327b 100644 --- a/backend/app/db/models.py +++ b/backend/app/db/models.py @@ -300,11 +300,37 @@ class Chunk(Base): # 而在此之前代码里没有任何地方能看出"这条是旧模型建的" embed_model: Mapped[str] = mapped_column(String(100), default="") embed_dim: Mapped[int] = mapped_column(Integer, default=0) + #: 分词后的 token 数。BM25 要按文档长度归一，存下来省得为算 avg_len 再扫一遍 + token_len: Mapped[int] = mapped_column(Integer, default=0) meta: Mapped[dict[str, Any]] = mapped_column(default=dict) document: Mapped[Document] = relationship(back_populates="chunks") +class ChunkTerm(Base): + """倒排表：词 → 包含它的片段。 + + 在此之前每次检索都 `select(Chunk)` 全量载入，再在 Python 里**重建**一遍 + BM25 倒排——几百段无感，上万段就是秒级，而那正是知识库开始有用的规模。 + + 没用 SQLite 的 FTS5：它的 trigram 分词器要求查询至少三个字符，中文里 + 「商家」「订单」「权限」这些两字词一个都召不回（实测）。而现有的分词器 + 本来就做中文单字 + 二元组，自己存一张倒排表既保住召回，也不把知识库 + 锁死在 SQLite 上。 + """ + + __tablename__ = "chunk_terms" + __table_args__ = (Index("ix_chunk_terms_lookup", "collection", "term"),) + + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + chunk_id: Mapped[str] = mapped_column( + ForeignKey("chunks.id", ondelete="CASCADE"), index=True + ) + collection: Mapped[str] = mapped_column(String(100), default="default") + term: Mapped[str] = mapped_column(String(64)) + tf: Mapped[int] = mapped_column(Integer, default=1) + + class Skill(Base, TimestampMixin): """可复用的方法论：一段指令 + 可选的示例和推荐工具。 diff --git a/backend/app/memory/embeddings.py b/backend/app/memory/embeddings.py index 24b111e2..0ab3a833 100644 --- a/backend/app/memory/embeddings.py +++ b/backend/app/memory/embeddings.py @@ -260,16 +260,21 @@ def hybrid_rank( query_vec: np.ndarray, *, alpha: float = 0.5, + keyword_scores: Sequence[float] | None = None, ) -> list[tuple[int, float, dict[str, float]]]: """向量相似度和 BM25 各归一化后加权合并。 纯向量会漏掉精确的专有名词，纯 BM25 又抓不住同义表达， 混合排序在本地小知识库上稳定优于任何一个单独使用。 + + keyword_scores 传进来就不再自己算 BM25。倒排索引那条路必须传：它按**全集** + 统计（N、avg_len、df）算好了分，而这里只拿到候选子集——在子集上重算， + df 和平均长度全是另一套数，排出来的名次和全表扫不一样。真踩过， + test_the_index_ranks_the_same_as_a_full_scan 就是为此写的。 """ if not texts: return [] - bm = BM25(texts) - kw_scores = bm.score(query) + kw_scores = list(keyword_scores) if keyword_scores is not None else BM25(texts).score(query) vec_scores = [cosine(query_vec, v) if v is not None else 0.0 for v in vectors] def _norm(xs: list[float]) -> list[float]: diff --git a/backend/app/memory/inverted.py b/backend/app/memory/inverted.py new file mode 100644 index 00000000..696c4363 --- /dev/null +++ b/backend/app/memory/inverted.py @@ -0,0 +1,131 @@ +"""倒排索引：把候选从"整个集合"收窄到"含有查询词的那些"。 + +在此之前每次检索都把集合里所有片段载入内存，再重建一遍 BM25 倒排。这不只是 +扫表——是对每一段重新分词、重新统计词频，每次查询都来一遍。 + +这里把倒排落到库里（chunk_terms），查询时只捞含有查询词的片段。BM25 的公式 +和参数与原来完全一致（k1=1.5, b=0.75，同一个分词器），所以排序结果不变—— +test_retrieval_quality 里有一条专门比对两条路径给出的名次是否一样。 +""" + +from __future__ import annotations + +import math +from collections import Counter, defaultdict +from typing import Any + +from sqlalchemy import delete, func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from app.db.models import Chunk, ChunkTerm +from app.memory.embeddings import _tokens + +_K1 = 1.5 +_B = 0.75 + + +def term_freqs(text: str) -> dict[str, int]: + """一段文本的词频。和 BM25 用同一个分词器，否则倒排和打分会对不上。""" + freqs: dict[str, int] = defaultdict(int) + for tok in _tokens(text): + # 超长 token 存不进 String(64)，而它们几乎都是噪音（长串 id、base64） + if len(tok) <= 64: + freqs[tok] += 1 + return dict(freqs) + + +async def index_chunk(session: AsyncSession, chunk: Chunk) -> None: + """给一个片段建倒排。调用方负责 commit。""" + freqs = term_freqs(chunk.content) + chunk.token_len = sum(freqs.values()) + for term, tf in freqs.items(): + session.add(ChunkTerm( + chunk_id=chunk.id, collection=chunk.collection, term=term, tf=tf)) + + +async def drop_chunk_terms(session: AsyncSession, chunk_ids: list[str]) -> None: + """删片段时一起删倒排。 + + 外键写了 ondelete=CASCADE，但 SQLite 的级联要连接级 PRAGMA 打开才生效 + （base.py 里是打开的）——索引一致性不该依赖一个开关，显式删。 + """ + if chunk_ids: + await session.execute( + delete(ChunkTerm).where(ChunkTerm.chunk_id.in_(chunk_ids))) + + +async def candidates( + session: AsyncSession, *, collection: str | None, query: str, limit: int, +) -> list[tuple[str, float]]: + """按 BM25 取候选，返回 [(chunk_id, 分数)]，已按分数降序。 + + 没有任何查询词命中时返回空——调用方据此决定是不是要退回全表。 + """ + # 按查询词的出现次数加权，不能去重。原来的 BM25.score 遍历的是带重复的 + # 词表：「导出订单最多多少行」里「多」出现两次，它的贡献就加两次。去重会让 + # 同一个查询在倒排和全表扫两条路上给出不同的分——实测顶部那条差了 2.32 + qcounts = Counter(t for t in _tokens(query) if len(t) <= 64) + terms = list(qcounts) + if not terms: + return [] + + stmt = select(ChunkTerm.chunk_id, ChunkTerm.term, ChunkTerm.tf).where( + ChunkTerm.term.in_(terms)) + if collection: + stmt = stmt.where(ChunkTerm.collection == collection) + rows = list(await session.execute(stmt)) + if not rows: + return [] + + # df 直接从命中里数：含有该词的片段数，正是 BM25 要的那个 df + df: dict[str, int] = defaultdict(int) + for _chunk_id, term, _tf in rows: + df[term] += 1 + + stat = select(func.count(Chunk.id), func.sum(Chunk.token_len)) + if collection: + stat = stat.where(Chunk.collection == collection) + total, total_len = (await session.execute(stat)).one() + total = int(total or 0) + if not total: + return [] + avg_len = float(total_len or 0) / total or 1.0 + + # 要 .all()：Result 自带 keys()，dict() 会当它是映射去下标取值， + # 抛 'ChunkedIteratorResult' object is not subscriptable + lengths = dict((await session.execute( + select(Chunk.id, Chunk.token_len).where( + Chunk.id.in_({r[0] for r in rows})))).all()) + + scores: dict[str, float] = defaultdict(float) + for chunk_id, term, tf in rows: + idf = math.log(1 + (total - df[term] + 0.5) / (df[term] + 0.5)) + length = lengths.get(chunk_id) or 1 + denom = tf + _K1 * (1 - _B + _B * length / avg_len) + scores[chunk_id] += qcounts[term] * idf * (tf * (_K1 + 1)) / denom + + ranked = sorted(scores.items(), key=lambda kv: kv[1], reverse=True) + return ranked[:limit] + + +async def missing_count(session: AsyncSession, collection: str | None = None) -> int: + """有多少片段还没建倒排。存量数据升级上来时就是全部。""" + stmt = select(func.count(Chunk.id)).where(Chunk.token_len == 0) + if collection: + stmt = stmt.where(Chunk.collection == collection) + return int((await session.execute(stmt)).scalar_one() or 0) + + +async def rebuild(session: AsyncSession, collection: str | None = None) -> int: + """重建倒排。存量片段升级后要跑一次，改了分词器也要。""" + stmt = select(Chunk) + if collection: + stmt = stmt.where(Chunk.collection == collection) + rows = list((await session.execute(stmt)).scalars()) + if not rows: + return 0 + await drop_chunk_terms(session, [r.id for r in rows]) + for row in rows: + await index_chunk(session, row) + await session.commit() + return len(rows) diff --git a/backend/app/memory/kb.py b/backend/app/memory/kb.py index 29b09dbb..a24f0db5 100644 --- a/backend/app/memory/kb.py +++ b/backend/app/memory/kb.py @@ -7,6 +7,7 @@ from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession from app.db.models import Chunk, Document +from app.memory import inverted from app.memory.embeddings import ( default_alpha, embed_text, embed_texts, embedder_dim, embedder_id, from_blob, hybrid_rank, to_blob, usable, @@ -90,21 +91,26 @@ async def ingest_document( await session.flush() pieces = chunk_text(content) + made: list[Chunk] = [] if pieces: vectors = await embed_texts(pieces) for i, (piece, vec) in enumerate(zip(pieces, vectors)): - session.add( - Chunk( - document_id=doc.id, - collection=collection, - ordinal=i, - content=piece, - embedding=to_blob(vec), - embed_model=embedder_id(), - embed_dim=embedder_dim(), - meta={"title": doc.title}, - ) + chunk = Chunk( + document_id=doc.id, + collection=collection, + ordinal=i, + content=piece, + embedding=to_blob(vec), + embed_model=embedder_id(), + embed_dim=embedder_dim(), + meta={"title": doc.title}, ) + session.add(chunk) + made.append(chunk) + # 先 flush 拿到 chunk.id，倒排行要引用它 + await session.flush() + for chunk in made: + await inverted.index_chunk(session, chunk) doc.chunk_count = len(pieces) await session.commit() await session.refresh(doc) @@ -134,19 +140,54 @@ async def search( if alpha is None: alpha = default_alpha() - stmt = select(Chunk) - if collection: - stmt = stmt.where(Chunk.collection == collection) - rows = list((await session.execute(stmt)).scalars()) + # 先用倒排把候选收窄。以前这里是把整个集合载入内存、再在 Python 里重建 + # 一遍 BM25 倒排——每次查询都重新分词一遍所有片段。 + # + # 多捞一些（limit 的 8 倍）：倒排只按关键词收，最终名次还要混进向量那一路， + # 卡得太紧会把"关键词一般但语义很近"的那些提前筛掉。 + rows: list[Chunk] = [] + kw_scores: list[float] | None = None + picked = await inverted.candidates( + session, collection=collection, query=query, limit=max(limit * 8, 50)) + if picked: + found = dict( + (c.id, c) for c in (await session.execute( + select(Chunk).where(Chunk.id.in_([cid for cid, _ in picked])))).scalars()) + pairs = [(found[cid], sc) for cid, sc in picked if cid in found] + rows = [c for c, _ in pairs] + # 倒排算好的分要带下去。在候选子集上重算 BM25，df 和平均长度是另一套数 + kw_scores = [sc for _, sc in pairs] + + if not rows: + # 一个查询词都没命中，或者这批片段还没建倒排（存量数据升级上来就是这样）。 + # 退回全表扫，慢但不会漏——宁可慢，也不要因为索引没建好就说"没搜到" + stmt = select(Chunk) + if collection: + stmt = stmt.where(Chunk.collection == collection) + rows = list((await session.execute(stmt)).scalars()) + kw_scores = None # 全表扫这条路照旧自己算 if not rows: return [] + # 按 id 定序再排。hybrid_rank 用的是稳定排序，分数并列时保留输入顺序—— + # 而倒排那条路和全表扫那条路喂进去的顺序天生不同，同样的查询会因为走了 + # 哪条路给出不同的名次。并列本来就该有个确定的先后，用 id 就够了 + if kw_scores is not None: + paired = sorted(zip(rows, kw_scores), key=lambda t: t[0].id) + rows = [r for r, _ in paired] + kw_scores = [s2 for _, s2 in paired] + else: + rows = sorted(rows, key=lambda r: r.id) + stale = [r for r in rows if not usable(r.embed_model, r.embed_dim)] if stale: if on_degrade: names = sorted({r.embed_model or "未知模型" for r in stale}) + # 报的是整个集合里要重建多少条，不是本次候选里有几条——候选被倒排 + # 收窄过，用它会把问题说小，而用户要照着做的恰恰是"重建多少" + total_stale = await stale_count(session, collection) on_degrade( - f"{collection or '全部集合'}里有 {len(stale)}/{len(rows)} 条向量由" + f"{collection or '全部集合'}里有 {total_stale} 条向量由" f"「{'、'.join(names)}」建立，和当前的「{embedder_id()}」对不上，" f"本次已退回纯关键词检索。重建索引后恢复。" ) @@ -163,6 +204,7 @@ async def search( vectors, query_vec, alpha=alpha, + keyword_scores=kw_scores, ) out: list[dict[str, Any]] = [] for idx, score, parts in ranked[:limit]: @@ -204,22 +246,33 @@ async def reindex(session: AsyncSession, collection: str | None = None) -> dict[ row.embed_model = model_id row.embed_dim = dim await session.commit() + # 倒排一起重建：两者都是"存量数据升级"的恢复手段，分成两个按钮只会漏点一个 + await inverted.rebuild(session, collection) return {"collection": collection, "reindexed": len(rows), "embedder": model_id} async def stale_count(session: AsyncSession, collection: str | None = None) -> int: - """有多少条向量和当前 embedder 对不上。知识库页据此提示要不要重建。""" - stmt = select(Chunk) + """有多少条向量和当前 embedder 对不上。知识库页据此提示要不要重建。 + + 用 SQL 数而不是把片段全载进来再在 Python 里过一遍——这个数在检索降级时 + 也要取，那条路径上不该再来一次全表扫。 + """ + from sqlalchemy import func, or_ + + stmt = select(func.count(Chunk.id)).where( + or_(Chunk.embed_model != embedder_id(), Chunk.embed_dim != embedder_dim())) if collection: stmt = stmt.where(Chunk.collection == collection) - rows = list((await session.execute(stmt)).scalars()) - return sum(1 for r in rows if not usable(r.embed_model, r.embed_dim)) + return int((await session.execute(stmt)).scalar_one() or 0) async def delete_document(session: AsyncSession, doc_id: str) -> bool: doc = await session.get(Document, doc_id) if not doc: return False + chunk_ids = list((await session.execute( + select(Chunk.id).where(Chunk.document_id == doc_id))).scalars()) + await inverted.drop_chunk_terms(session, chunk_ids) await session.delete(doc) await session.commit() return True diff --git a/backend/app/memory/parsing.py b/backend/app/memory/parsing.py new file mode 100644 index 00000000..d6d85e1c --- /dev/null +++ b/backend/app/memory/parsing.py @@ -0,0 +1,120 @@ +"""把上传的文件变成文本。 + +在此之前 upload 只接受 UTF-8：PDF 传上去直接 415，而 mime 字段记了却从不 +参与任何判断。知识库里最常见的恰恰是 PDF 和 Word。 + +PDF/DOCX 的解析库走可选依赖（pyproject 的 [docs]），和 [db]、[microvm] 一样 +——它们加起来几十兆，而只用纯文本的人不该为此付出安装时间。没装时给的报错 +要能照着做，不是一句 ImportError。 + +HTML 用已经在依赖里的 beautifulsoup4，不额外引入东西。 +""" + +from __future__ import annotations + +import io +import re + +#: 认得出的扩展名 → 处理方式。mime 不可靠（浏览器给 .md 报 application/octet-stream） +_PDF = {".pdf"} +_DOCX = {".docx"} +_HTML = {".html", ".htm"} + + +class UnsupportedDocument(ValueError): + """认不出或解析不了。message 直接给用户看。""" + + +def _ext(filename: str) -> str: + idx = filename.rfind(".") + return filename[idx:].lower() if idx >= 0 else "" + + +def _from_pdf(raw: bytes) -> str: + try: + from pypdf import PdfReader + except ImportError as e: + raise UnsupportedDocument( + "要读 PDF 得先装解析库：pip install 'agentlab-backend[docs]'" + ) from e + try: + reader = PdfReader(io.BytesIO(raw)) + except Exception as e: # noqa: BLE001 + raise UnsupportedDocument(f"这个 PDF 读不开：{type(e).__name__}: {e}") from e + + # 按页拼，页与页之间留空行——切块是按段落切的，页边界正好是天然的段落边界 + pages = [] + for i, page in enumerate(reader.pages): + try: + pages.append(page.extract_text() or "") + except Exception: # noqa: BLE001 - 单页坏了不该毁掉整份文档 + pages.append(f"（第 {i + 1} 页解析失败）") + # PDF 抽出来的每行都拖着排版留下的填充空格，不清掉会混进切块和词频统计 + cleaned = [] + for page in pages: + lines = [ln.rstrip() for ln in page.splitlines()] + body = "\n".join(ln for ln in lines if ln.strip()) + if body: + cleaned.append(body) + text = "\n\n".join(cleaned) + if not text.strip(): + raise UnsupportedDocument( + "这个 PDF 里没有可提取的文字——多半是扫描件。" + "扫描件要先过 OCR，这里不做图像识别。" + ) + return text + + +def _from_docx(raw: bytes) -> str: + try: + import docx + except ImportError as e: + raise UnsupportedDocument( + "要读 Word 得先装解析库：pip install 'agentlab-backend[docs]'" + ) from e + try: + doc = docx.Document(io.BytesIO(raw)) + except Exception as e: # noqa: BLE001 + raise UnsupportedDocument(f"这个 Word 文档读不开：{type(e).__name__}: {e}") from e + + parts = [p.text.strip() for p in doc.paragraphs if p.text.strip()] + # 表格也要：合同、口径说明这类文档的关键信息常常整个在表里， + # 只取段落等于把它们全丢了 + for table in doc.tables: + for row in table.rows: + cells = [c.text.strip() for c in row.cells if c.text.strip()] + if cells: + parts.append(" | ".join(cells)) + text = "\n\n".join(parts) + if not text.strip(): + raise UnsupportedDocument("这个 Word 文档里没有文字内容") + return text + + +def _from_html(raw: bytes) -> str: + from bs4 import BeautifulSoup + + soup = BeautifulSoup(raw.decode("utf-8", errors="replace"), "html.parser") + for tag in soup(["script", "style", "nav", "footer"]): + tag.decompose() + text = soup.get_text("\n") + # 连续空行压成一个：HTML 转出来到处是空行，会把切块的段落判断搅乱 + return re.sub(r"\n{3,}", "\n\n", text).strip() + + +def extract(raw: bytes, filename: str, mime: str = "") -> str: + """文件字节 → 文本。认不出就按 UTF-8 文本试。""" + ext = _ext(filename) + if ext in _PDF or mime == "application/pdf": + return _from_pdf(raw) + if ext in _DOCX or "wordprocessingml" in mime: + return _from_docx(raw) + if ext in _HTML or mime == "text/html": + return _from_html(raw) + try: + return raw.decode("utf-8") + except UnicodeDecodeError as e: + raise UnsupportedDocument( + f"读不了 {filename or '这个文件'}：它不是 UTF-8 文本，" + f"也不是认得出的 PDF / Word / HTML。" + ) from e diff --git a/backend/pyproject.toml b/backend/pyproject.toml index d06f16ef..144adca7 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -33,6 +33,9 @@ dependencies = [ microvm = ["microsandbox>=0.6"] # 测试：conda 环境默认不装，跑测试前 pip install 'agentlab-backend[dev]' dev = ["pytest>=8.0", "pytest-asyncio>=0.24"] +# 文档解析。PDF 和 Word 的解析库加起来几十兆，只用纯文本的人不该为此 +# 付出安装时间——和 microvm、db 一样做成可选。HTML 用已在依赖里的 bs4。 +docs = ["pypdf>=5.0", "python-docx>=1.1"] # 数据库接入。三个驱动都是 asyncio 原生且免装数据库客户端—— # oracledb 走 thin 模式，省掉了 Oracle Instant Client 这个历来最劝退的一步。 db = ["aiomysql>=0.2", "oracledb>=2.0", "asyncpg>=0.30"] diff --git a/backend/tests/test_kb.py b/backend/tests/test_kb.py index 5eafd6ab..f1ebaa78 100644 --- a/backend/tests/test_kb.py +++ b/backend/tests/test_kb.py @@ -131,7 +131,7 @@ async def test_search_degrades_instead_of_crashing_after_a_model_switch() -> Non assert hits, "退回关键词也要能搜到东西，而不是空手而归" assert notes, "少了一半能力却不吭声，用户只会觉得最近搜得不准" - assert "重建索引" in notes[0] and "2/2" in notes[0] + assert "重建索引" in notes[0] and "2 条" in notes[0], notes[0] async def test_reindex_restores_hybrid_search() -> None: @@ -172,3 +172,76 @@ def test_choosing_a_remote_embedder_that_cannot_be_built_says_so(monkeypatch) -> def test_the_local_embedder_is_always_available() -> None: emb.configure("local") assert emb.embedder_id() == "local-hashing" and emb.embedder_dim() == 512 + + +# -------------------------------------------------------------------------- +# 文档解析 +# -------------------------------------------------------------------------- +# +# 以前 upload 只接受 UTF-8，PDF 传上去直接 415；而 mime 字段记了却从不参与 +# 任何判断。知识库里最常见的恰恰是 PDF 和 Word。 + + +def test_plain_text_still_goes_straight_through() -> None: + from app.memory.parsing import extract + + assert extract("就是一段文本".encode(), "note.txt") == "就是一段文本" + assert extract("# 标题".encode(), "readme.md") == "# 标题" + + +def test_html_is_stripped_to_text() -> None: + from app.memory.parsing import extract + + html = """ +

权限说明

只有管理员能改角色。

+ """.encode() + out = extract(html, "doc.html") + assert "权限说明" in out and "只有管理员能改角色" in out + # 脚本样式导航页脚都不是正文，留着只会污染检索 + assert "alert" not in out and "color:red" not in out + assert "导航" not in out and "版权" not in out + + +def test_an_unreadable_binary_says_what_to_do() -> None: + """报错要能照着做，不是一句 UnicodeDecodeError。""" + from app.memory.parsing import UnsupportedDocument, extract + + with pytest.raises(UnsupportedDocument) as e: + extract(b"\x00\x01\x02\xff\xfe", "mystery.bin") + assert "UTF-8" in str(e.value) and "PDF" in str(e.value) + + +def test_a_missing_parser_tells_you_which_package(monkeypatch) -> None: + """没装可选依赖时给的必须是"装哪个包"，而不是 ImportError。""" + import builtins + + from app.memory.parsing import UnsupportedDocument, extract + + real_import = builtins.__import__ + + def _no_pypdf(name, *a, **kw): + if name == "pypdf": + raise ImportError("no pypdf") + return real_import(name, *a, **kw) + + monkeypatch.setattr(builtins, "__import__", _no_pypdf) + with pytest.raises(UnsupportedDocument) as e: + extract(b"%PDF-1.4 fake", "doc.pdf") + assert "agentlab-backend[docs]" in str(e.value) + + +def test_a_scanned_pdf_is_called_out_rather_than_ingested_empty() -> None: + """扫描件提不出文字。悄悄存一份空文档，用户会以为传成功了、却永远搜不到。""" + pytest.importorskip("pypdf") + import pypdf + + from app.memory.parsing import UnsupportedDocument, extract + + writer = pypdf.PdfWriter() + writer.add_blank_page(width=200, height=200) + buf = __import__("io").BytesIO() + writer.write(buf) + + with pytest.raises(UnsupportedDocument) as e: + extract(buf.getvalue(), "scan.pdf") + assert "扫描件" in str(e.value) diff --git a/backend/tests/test_retrieval_quality.py b/backend/tests/test_retrieval_quality.py index 4072ff8d..fcb387cf 100644 --- a/backend/tests/test_retrieval_quality.py +++ b/backend/tests/test_retrieval_quality.py @@ -293,3 +293,93 @@ async def _measure_default(k: int = 5) -> dict[str, float]: rr.append(1.0 / rank if rank else 0.0) n = len(hit1) return {"all_top1": sum(hit1) / n, "all_hit": sum(hit) / n, "all_rr": sum(rr) / n} + + +# -------------------------------------------------------------------------- +# 倒排索引：必须和全表扫给出同一个名次 +# -------------------------------------------------------------------------- +# +# 这是换索引最要紧的一条。一个更快但排序不同的检索，等于悄悄改了所有人的 +# 检索结果——而且因为"看起来还是有结果"，没人会发现。 + + +async def _brute_force(query: str, k: int) -> list[str]: + """绕开倒排，按原来的方式全表扫 + 现建 BM25。""" + from sqlalchemy import select + + from app.db.models import Chunk + from app.memory.embeddings import embed_text, from_blob, hybrid_rank + + async with SessionLocal() as session: + rows = list((await session.execute( + select(Chunk).where(Chunk.collection == COLLECTION))).scalars()) + # 和 kb.search 一样按 id 定序：分数并列时排序是稳定的，输入顺序不同 + # 就会给出不同的名次。参照实现也得是确定的，否则这条比对时灵时不灵 + rows.sort(key=lambda r: r.id) + qv = await embed_text(query) + ranked = hybrid_rank(query, [r.content for r in rows], + [from_blob(r.embedding, emb.embedder_dim()) for r in rows], + qv, alpha=emb.default_alpha()) + return [rows[i].meta.get("title", "") for i, _s, _p in ranked[:k]] + + +async def test_the_index_ranks_the_same_as_a_full_scan() -> None: + """快是次要的，一致才是前提。""" + async with SessionLocal() as session: + for question, _want, _kind in QUESTIONS: + fast = [h["title"] for h in await kb.search( + session, collection=COLLECTION, query=question, limit=5)] + slow = await _brute_force(question, 5) + assert fast == slow, f"「{question}」倒排给的是 {fast}，全表扫给的是 {slow}" + + +async def test_the_index_actually_narrows_the_candidates() -> None: + """不收窄的话这一整套就白做了——查询词只该捞出含有它们的那些片段。""" + from app.memory import inverted + + async with SessionLocal() as session: + picked = await inverted.candidates( + session, collection=COLLECTION, query="会话多久过期", limit=100) + assert picked, "倒排一条都没捞到" + assert len(picked) < len(DOCS), ( + f"捞了 {len(picked)} 条，语料才 {len(DOCS)} 篇——等于没收窄") + + +async def test_chunks_without_an_index_fall_back_to_the_full_scan() -> None: + """存量数据升级上来时倒排是空的。宁可慢，也不要因为索引没建好就说"没搜到"。""" + from sqlalchemy import select + + from app.db.models import Chunk + from app.memory import inverted + + async with SessionLocal() as session: + ids = list((await session.execute( + select(Chunk.id).where(Chunk.collection == COLLECTION))).scalars()) + await inverted.drop_chunk_terms(session, ids) + await session.commit() + + hits = await kb.search(session, collection=COLLECTION, query="会话多久过期", limit=3) + assert hits, "倒排被清空后不该搜不到东西" + assert hits[0]["title"] == "登录与会话" + + # 重建回来，别影响后面的用例 + assert await inverted.rebuild(session, COLLECTION) == len(DOCS) + + +async def test_deleting_a_document_takes_its_index_rows() -> None: + """倒排不跟着删的话，它会一直把已经不存在的片段捞出来。""" + from sqlalchemy import func, select + + from app.db.models import ChunkTerm, Document + + async with SessionLocal() as session: + doc = await kb.ingest_document( + session, collection="idx_del", title="临时", content="一段会被删掉的内容。") + before = int((await session.execute(select(func.count(ChunkTerm.id)).where( + ChunkTerm.collection == "idx_del"))).scalar_one()) + assert before > 0 + + await kb.delete_document(session, doc.id) + after = int((await session.execute(select(func.count(ChunkTerm.id)).where( + ChunkTerm.collection == "idx_del"))).scalar_one()) + assert after == 0, f"文档删了但倒排还留着 {after} 行"