"""知识库:切块、混合检索、换 embedder 之后还能不能用。 这三个模块(kb.py / embeddings.py / memory/store.py,共 526 行检索逻辑) 在此之前**一条测试都没有**。于是有两件事一直没人发现: - 换 embedding 模型会让检索整体 500:库里是 512 维的本地向量、查询是 1536 维 的远端向量,cosine 直接 `shapes (1536,) and (512,) not aligned` - 存量向量属于哪个模型,代码里没有任何地方记着,所以也无从降级 现在这两件事都有守的了。 """ from __future__ import annotations import numpy as np import pytest from app.db.base import SessionLocal from app.db.models import Chunk from app.memory import embeddings as emb from app.memory import kb from app.memory.embeddings import hybrid_rank, to_blob @pytest.fixture(autouse=True) def _local_embedder(): """每个用例都从本地 embedder 起步,别让上一个用例的 configure 漏过来。""" emb.configure("local") yield emb.configure("local") @pytest.fixture async def client(): from httpx import ASGITransport, AsyncClient from app.main import app transport = ASGITransport(app=app) async with AsyncClient(transport=transport, base_url="http://test") as c: yield c # -------------------------------------------------------------------------- # 切块 # -------------------------------------------------------------------------- def test_short_paragraphs_stay_whole() -> None: blocks = kb.chunk_text("第一段。\n\n第二段。\n\n第三段。") assert blocks and all("第" in b for b in blocks) def test_a_long_paragraph_is_cut_on_sentence_boundaries() -> None: """命中之后读起来要是完整的句子,不能拦腰截断。""" text = "。".join(f"这是第{i}句话,讲了一点内容" for i in range(120)) + "。" blocks = kb.chunk_text(text, target=200, overlap=0) assert len(blocks) > 1 # 除了最后一块,其余都该收在句号上 assert all(b.rstrip().endswith("。") for b in blocks[:-1]) def test_a_sentence_longer_than_the_target_is_hard_cut() -> None: """没有标点可切时也不能无限膨胀,硬切是最后的兜底。""" blocks = kb.chunk_text("啊" * 2000, target=300, overlap=0) assert len(blocks) > 1 and max(len(b) for b in blocks) <= 300 def test_empty_input_yields_nothing() -> None: assert kb.chunk_text("") == [] and kb.chunk_text(" \n\n ") == [] # -------------------------------------------------------------------------- # hybrid_rank 的归一化边界 # -------------------------------------------------------------------------- def test_a_single_candidate_is_not_scored_zero() -> None: """min-max 归一化在只有一条候选时 hi==lo。 返回 0 的话它的合并分恒为 0,会被 min_score 过滤掉——于是新建集合写入 第一条之后永远搜不到,而"先写一条、下一轮再检索"正是 agent 最常见的用法。 embeddings.py 里为此写了一段注释,但一直没有测试守着。 """ vec = np.ones(4, dtype=np.float32) ranked = hybrid_rank("查询", ["只有这一条"], [vec], vec) assert len(ranked) == 1 and ranked[0][1] > 0 # -------------------------------------------------------------------------- # 换 embedder:以前这里直接 500 # -------------------------------------------------------------------------- async def _seed(collection: str, texts: list[str]) -> None: async with SessionLocal() as session: for t in texts: await kb.ingest_document( session, collection=collection, title=t[:10], content=t) def test_stale_vectors_are_recognised() -> None: """存量数据那两列是空的,一律当成对不上——保守但不会拿旧坐标冒充新空间。""" assert emb.usable("", 0) is False assert emb.usable("openai:text-embedding-3-small", 1536) is False assert emb.usable(emb.embedder_id(), emb.embedder_dim()) is True async def test_ingest_records_who_built_the_vector() -> None: await _seed("kbtest_mark", ["平台管理员一共有六个人。"]) async with SessionLocal() as session: from sqlalchemy import select row = (await session.execute( select(Chunk).where(Chunk.collection == "kbtest_mark"))).scalars().first() assert row.embed_model == emb.embedder_id() assert row.embed_dim == emb.embedder_dim() async def test_search_degrades_instead_of_crashing_after_a_model_switch() -> None: """整件事的核心。 以前:cosine(1536 维查询, 512 维存量) → ValueError → 检索 500。 现在:整体退回纯关键词,并把原因说出来。 """ await _seed("kbtest_switch", ["平台管理员一共有六个人。", "商家账户共五个。"]) # 把库里的向量伪装成别的模型建的(维度也不同) async with SessionLocal() as session: from sqlalchemy import select rows = list((await session.execute( select(Chunk).where(Chunk.collection == "kbtest_switch"))).scalars()) for r in rows: r.embedding = to_blob(np.ones(1536, dtype=np.float32)) r.embed_model = "openai:text-embedding-3-small" r.embed_dim = 1536 await session.commit() notes: list[str] = [] async with SessionLocal() as session: hits = await kb.search(session, collection="kbtest_switch", query="管理员", limit=5, on_degrade=notes.append) assert hits, "退回关键词也要能搜到东西,而不是空手而归" assert notes, "少了一半能力却不吭声,用户只会觉得最近搜得不准" assert "重建索引" in notes[0] and "2 条" in notes[0], notes[0] async def test_reindex_restores_hybrid_search() -> None: await _seed("kbtest_reindex", ["平台管理员一共有六个人。"]) async with SessionLocal() as session: from sqlalchemy import select row = (await session.execute( select(Chunk).where(Chunk.collection == "kbtest_reindex"))).scalars().first() row.embed_model, row.embed_dim = "别的模型", 99 await session.commit() async with SessionLocal() as session: assert await kb.stale_count(session, "kbtest_reindex") == 1 out = await kb.reindex(session, "kbtest_reindex") assert out["reindexed"] == 1 assert await kb.stale_count(session, "kbtest_reindex") == 0 notes: list[str] = [] async with SessionLocal() as session: hits = await kb.search(session, collection="kbtest_reindex", query="管理员", on_degrade=notes.append) assert hits and not notes, "重建之后不该再降级" # -------------------------------------------------------------------------- # 配置 embedder # -------------------------------------------------------------------------- def test_choosing_a_remote_embedder_that_cannot_be_built_says_so(monkeypatch) -> None: """选了远端却静默退回本地哈希向量,用户会以为自己在用语义检索。""" from langchain_openai import OpenAIEmbeddings monkeypatch.delenv("OPENAI_API_KEY", raising=False) # 建 embedder 时会真调一次接口探维度:单测不走网络(以前真连 api.openai.com, # 代理慢的时候这一条要等将近 4 分钟),直接给一个鉴权失败 def refuse(self, text): raise RuntimeError("Error code: 401 - invalid api key") monkeypatch.setattr(OpenAIEmbeddings, "embed_query", refuse) with pytest.raises(emb.EmbedderUnavailable) as e: emb.configure("openai", "text-embedding-3-small") assert "OPENAI_API_KEY" in str(e.value) 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) # -------------------------------------------------------------------------- # 长期记忆:和知识库同一个坑,以前只有知识库有恢复手段 # -------------------------------------------------------------------------- async def _memory(scope: str, content: str): from app.memory import store async with SessionLocal() as session: return await store.remember(session, scope=scope, content=content, dedupe=False) async def test_memory_records_who_built_the_vector() -> None: from app.db.models import MemoryItem item = await _memory("memtest_mark", "后端跑在 8001 端口。") async with SessionLocal() as session: row = await session.get(MemoryItem, item.id) assert row.embed_model == emb.embedder_id() and row.embed_dim == emb.embedder_dim() async def test_memory_recall_degrades_out_loud() -> None: """以前退得悄无声息。知识库那边至少会发一条 warn,记忆这边只是"最近想不起 事",而人根本不会想到要去重建索引。""" from sqlalchemy import select from app.db.models import MemoryItem from app.memory import store await _memory("memtest_stale", "公司年会定在十二月二十号。") async with SessionLocal() as session: row = (await session.execute(select(MemoryItem).where( MemoryItem.scope == "memtest_stale"))).scalars().first() row.embed_model, row.embed_dim = "openai:text-embedding-3-small", 1536 row.embedding = to_blob(np.ones(1536, dtype=np.float32)) await session.commit() notes: list[str] = [] async with SessionLocal() as session: hits = await store.recall(session, scope="memtest_stale", query="年会", on_degrade=notes.append) assert hits, "退回关键词也要能召回,而不是空手" assert notes and "重建索引" in notes[0], notes async def test_memory_reindex_restores_it() -> None: from sqlalchemy import select from app.db.models import MemoryItem from app.memory import store await _memory("memtest_fix", "张三的邮箱是 zhangsan@example.com。") async with SessionLocal() as session: row = (await session.execute(select(MemoryItem).where( MemoryItem.scope == "memtest_fix"))).scalars().first() row.embed_model, row.embed_dim = "别的模型", 99 await session.commit() assert await store.stale_count(session, "memtest_fix") == 1 assert await store.reindex(session, "memtest_fix") == 1 assert await store.stale_count(session, "memtest_fix") == 0 notes: list[str] = [] async with SessionLocal() as session: hits = await store.recall(session, scope="memtest_fix", query="邮箱", on_degrade=notes.append) assert hits and not notes, "重建之后不该再降级" async def test_one_rebuild_covers_both_stores() -> None: """知识库和记忆共用一个 embedder,换模型时一起失效。 分成两个按钮的结果是点了一个、以为好了,另一个还在悄悄退回关键词。 """ from httpx import ASGITransport, AsyncClient from sqlalchemy import select from app.db.models import Chunk, MemoryItem from app.main import app from app.memory import store await _seed("both_fix", ["平台管理员有六个。"]) await _memory("both_mem", "商家账户共五个。") async with SessionLocal() as session: for model_cls, scope_col, scope_val in ( (Chunk, Chunk.collection, "both_fix"), (MemoryItem, MemoryItem.scope, "both_mem") ): for row in (await session.execute( select(model_cls).where(scope_col == scope_val))).scalars(): row.embed_model, row.embed_dim = "旧模型", 7 await session.commit() transport = ASGITransport(app=app) async with AsyncClient(transport=transport, base_url="http://test") as c: status = (await c.get("/api/kb/embedding")).json() assert status["stale_chunks"] >= 1 and status["stale_memories"] >= 1 out = (await c.post("/api/kb/reindex")).json() assert out["reindexed"] >= 1 assert out["memories_reindexed"] >= 1, "重建只管了知识库,记忆还留在旧模型上" after = (await c.get("/api/kb/embedding")).json() assert after["stale_chunks"] == 0 and after["stale_memories"] == 0 # -------------------------------------------------------------------------- # 表格横跨多块时还读不读得懂 # -------------------------------------------------------------------------- def _table_doc(rows: int = 60) -> str: """一份表格长到必然被切开的文档。""" body = ["门店 | 城市 | 销售额 | 同比"] body += [f"门店{i:02d} | 上海 | {100000 + i * 3137} | {i - 10}%" for i in range(rows)] return "季度报告\n\n" + "\n\n".join(body) def test_a_split_table_keeps_its_header_in_every_chunk() -> None: """一张几十行的表会横跨好几块,除了第一块其余都只有数据行。 「上海 | 172151 | 13%」这样的片段召回出来,模型没法知道这几个数是什么。 表头才几十个字,每块带一份换来的是这块单独拿出来也读得懂。 """ chunks = kb.chunk_text(_table_doc()) assert len(chunks) > 1, "语料不够长,没切开就测不到这件事" for i, c in enumerate(chunks[1:], 1): assert "门店 | 城市 | 销售额 | 同比" in c, f"块{i} 丢了表头" def test_the_header_leads_the_chunk_not_buried_in_it() -> None: """表头要在最前面。 第一版写成 tail+header+block,表头被夹在重叠的那几行后面成了中间一行, 块的开头仍然是一串没有主语的数——查 has_header 是 True,看起来修好了, 实际没有。 """ chunks = kb.chunk_text(_table_doc()) for i, c in enumerate(chunks[1:], 1): assert c.splitlines()[0].startswith("门店 | 城市"), ( f"块{i} 的第一行是「{c.splitlines()[0][:30]}」,表头没领头") def test_overlap_never_severs_a_row() -> None: """重叠原来按字符切(buf[-120:]),会从行中间断开。 实测切出来的块开头是「| 上海 | 172151 | 13%」——门店名被切没了, 剩一列没有主语的数字。 """ for c in kb.chunk_text(_table_doc()): first = c.splitlines()[0].lstrip() assert not first.startswith("|"), f"首行被腰斩:{first[:40]}" def test_plain_prose_is_unaffected_by_the_table_handling() -> None: """表格那套不能把普通文档搞坏。""" text = "\n\n".join(f"第{i}段。" + "内容" * 60 for i in range(8)) chunks = kb.chunk_text(text) assert len(chunks) > 1 assert not any(c.startswith("门店") for c in chunks) # 正文一个字都不能少 assert all(f"第{i}段。" in "".join(chunks) for i in range(8)) async def test_a_word_document_with_a_table_survives_the_whole_pipeline() -> None: """解析 → 切块 → 检索,端到端拿回来的片段要能读懂。""" docx = pytest.importorskip("docx") import io from app.memory.parsing import extract d = docx.Document() d.add_heading("门店运营手册", 0) d.add_paragraph("平台管理员拥有 platform 级别角色。") t = d.add_table(rows=1, cols=4) for i, h in enumerate(("门店", "城市", "销售额", "同比")): t.cell(0, i).text = h for i in range(60): row = t.add_row() for j, v in enumerate((f"门店{i:02d}", "上海", str(100000 + i * 3137), f"{i - 10}%")): row.cells[j].text = v buf = io.BytesIO() d.save(buf) text = extract(buf.getvalue(), "手册.docx") assert "平台管理员" in text and "门店00 | 上海" in text async with SessionLocal() as session: await kb.ingest_document(session, collection="docx_e2e", title="手册", content=text) hits = await kb.search(session, collection="docx_e2e", query="门店 销售额", limit=3) assert hits # 命中的片段里必须能看出这些数是什么 assert any("门店 | 城市 | 销售额" in h["content"] for h in hits), ( "召回的片段全都没有表头,拿回去也读不懂") # -------------------------------------------------------------------------- # 自定义 embedding 端点 # -------------------------------------------------------------------------- def test_dimension_is_probed_not_guessed(monkeypatch) -> None: """维度必须问出来。 原来是按模型名猜:名字里有 3-large 就 3072,否则一律 1536。本地部署的 模型五花八门——bge-m3 是 1024、Qwen3-Embedding-4B 是 2560——猜错的后果是 所有存量向量都被判成"和当前 embedder 对不上",检索整体退回关键词。 """ from app.memory import embeddings as e probed: list[str] = [] class _Fake: def __init__(self, **kw): self.kw = kw def embed_query(self, text): probed.append(text) return [0.0] * 2560 # 假装是个 2560 维的本地模型 import langchain_openai monkeypatch.setattr(langchain_openai, "OpenAIEmbeddings", _Fake) emb_obj = e.OpenAIEmbedder(model="my-local-model", base_url="http://127.0.0.1:1234/v1") assert emb_obj.dim == 2560, "维度没去问,还在按名字猜" assert probed, "根本没探测" def test_a_custom_endpoint_gets_plain_strings(monkeypatch) -> None: """langchain 默认先分词、把 token 数组发过去。 OpenAI 官方认,但 LM Studio 这类本地服务只收字符串,会直接 400: 'input' field must be a string or an array of strings """ from app.memory import embeddings as e captured: dict = {} class _Fake: def __init__(self, **kw): captured.update(kw) def embed_query(self, _t): return [0.0] * 8 import langchain_openai monkeypatch.setattr(langchain_openai, "OpenAIEmbeddings", _Fake) e.OpenAIEmbedder(model="m", base_url="http://127.0.0.1:1234/v1") assert captured.get("check_embedding_ctx_length") is False assert captured.get("base_url") == "http://127.0.0.1:1234/v1" # 批量上限:远端接口单次有上限,不分块的话大文档必挂 assert captured.get("chunk_size") captured.clear() e.OpenAIEmbedder(model="text-embedding-3-small") # 官方接口走默认行为 assert "check_embedding_ctx_length" not in captured def test_a_broken_endpoint_names_where_it_tried(monkeypatch) -> None: """报错要说清楚是哪个地址上的哪个模型,否则本地服务没起都不知道去哪查。""" from app.memory import embeddings as e class _Boom: def __init__(self, **kw): raise RuntimeError("connection refused") import langchain_openai monkeypatch.setattr(langchain_openai, "OpenAIEmbeddings", _Boom) with pytest.raises(e.EmbedderUnavailable) as err: e.configure("openai", "qwen3-embedding-4b", "http://127.0.0.1:1234/v1") assert "127.0.0.1:1234" in str(err.value) and "qwen3-embedding-4b" in str(err.value) async def test_the_probe_endpoint_lists_models(monkeypatch) -> None: """模型名手填太容易错:LM Studio 里叫 text-embedding-qwen3-embedding-4b, 不是 Qwen3-Embedding-4B。填错的表现是切换时 404,人还以为是服务没起。""" import httpx from httpx import ASGITransport, AsyncClient from app.main import app class _Resp: def json(self): return {"data": [{"id": "text-embedding-qwen3-embedding-4b"}, {"id": "other"}]} class _Client: def __init__(self, **kw): ... async def __aenter__(self): return self async def __aexit__(self, *a): return False async def get(self, _url): return _Resp() monkeypatch.setattr(httpx, "AsyncClient", _Client) transport = ASGITransport(app=app) async with AsyncClient(transport=transport, base_url="http://test") as c: out = (await c.get("/api/kb/embedding/probe", params={"base_url": "http://127.0.0.1:1234/v1"})).json() assert "text-embedding-qwen3-embedding-4b" in out["models"] async def test_a_dead_endpoint_says_which_address(monkeypatch) -> None: import httpx from httpx import ASGITransport, AsyncClient from app.main import app class _Client: def __init__(self, **kw): ... async def __aenter__(self): return self async def __aexit__(self, *a): return False async def get(self, _url): raise ConnectionError("refused") monkeypatch.setattr(httpx, "AsyncClient", _Client) transport = ASGITransport(app=app) async with AsyncClient(transport=transport, base_url="http://test") as c: r = await c.get("/api/kb/embedding/probe", params={"base_url": "http://127.0.0.1:9999/v1"}) assert r.status_code == 400 assert "127.0.0.1:9999" in r.json()["detail"] async def test_a_saved_model_that_did_not_connect_is_reported(monkeypatch) -> None: """配了本机 LM Studio 上的模型,服务启动时它还没起:进程一直用本地哈希。 以前状态接口只报 local-hashing,页面按"配的是不是 local"判断,于是既不提示 没有语义能力,还高亮"重建索引"——那些"对不上"的向量正是那个模型建好的, 一点就拿哈希覆盖掉。退回要报出来、原因要带上,重新连上之后提示要消失; 明确选了本地哈希的不算退回。 """ import langchain_openai from httpx import ASGITransport, AsyncClient from app.db.models import Setting from app.main import app class _Down: def __init__(self, **kw): raise RuntimeError("connection refused") class _Up: def __init__(self, **kw): pass def embed_query(self, text): return [0.1] * 8 # 模块级缓存和库里的配置都是整个测试会话共用的,结束时还原 monkeypatch.setattr(emb, "_cached", None) monkeypatch.setattr(emb, "_unavailable", None) saved = {"kind": "openai", "model": "text-embedding-qwen3-embedding-4b", "base_url": "http://127.0.0.1:1234/v1"} async with SessionLocal() as session: row = await session.get(Setting, emb.EMBEDDING_SETTING_KEY) before = dict(row.value) if row else None if row: row.value = saved else: session.add(Setting(key=emb.EMBEDDING_SETTING_KEY, value=saved)) await session.commit() try: monkeypatch.setattr(langchain_openai, "OpenAIEmbeddings", _Down) async with SessionLocal() as session: with pytest.raises(emb.EmbedderUnavailable): await emb.load_setting(session) transport = ASGITransport(app=app) async with AsyncClient(transport=transport, base_url="http://test") as c: down = (await c.get("/api/kb/embedding")).json() assert down["embedder"] == "local-hashing" and down["has_semantics"] is False assert down["fallback"] is True, "配了语义模型却在用哈希,状态里得说出来" assert "127.0.0.1:1234" in down["fallback_reason"], down["fallback_reason"] # 服务起来之后点「重新连接」:发的就是保存着的那份配置 monkeypatch.setattr(langchain_openai, "OpenAIEmbeddings", _Up) assert (await c.put("/api/kb/embedding", json=saved)).status_code == 200 up = (await c.get("/api/kb/embedding")).json() assert up["has_semantics"] is True assert up["fallback"] is False and up["fallback_reason"] == "" await c.put("/api/kb/embedding", json={"kind": "local"}) local = (await c.get("/api/kb/embedding")).json() assert local["has_semantics"] is False and local["fallback"] is False finally: async with SessionLocal() as session: row = await session.get(Setting, emb.EMBEDDING_SETTING_KEY) if before is None: await session.delete(row) else: row.value = before await session.commit() async def test_the_endpoint_is_remembered_across_restarts(client=None) -> None: """base_url 不落库的话,重启后会悄悄退回本地哈希——而库里的向量是 上一个模型建的,检索会整体退回关键词,用户只看到"重启之后搜得不准了"。""" from httpx import ASGITransport, AsyncClient from app.db.models import Setting from app.main import app from app.memory import embeddings as e transport = ASGITransport(app=app) async with AsyncClient(transport=transport, base_url="http://test") as c: await c.put("/api/kb/embedding", json={"kind": "local"}) status = (await c.get("/api/kb/embedding")).json() assert "base_url" in status, "状态里不回显 base_url,界面刷新后就不知道接的是哪" async with SessionLocal() as session: row = await session.get(Setting, e.EMBEDDING_SETTING_KEY) assert row is not None and "base_url" in row.value # -------------------------------------------------------------------------- # 大文件:分批、后台、大小防护 # -------------------------------------------------------------------------- async def test_embedding_is_batched_not_one_giant_call(monkeypatch) -> None: """原来是 embed_texts(pieces) 一次全发。 本地哈希只是慢,换成远端模型就是必挂:一个 10MB 文档切出近八千段, 远端接口单次批量有上限(OpenAI 是 2048),一次请求塞不下。 reindex 那边早就分批了,ingest 这条一直漏着。 """ from app.memory import kb as kb_mod sizes: list[int] = [] real = kb_mod.embed_texts async def spy(texts): sizes.append(len(texts)) return await real(texts) monkeypatch.setattr(kb_mod, "embed_texts", spy) text = "\n\n".join(f"第{i}段内容,讲了一些东西。" * 12 for i in range(400)) async with SessionLocal() as session: await kb.ingest_document(session, collection="batch_test", title="大文档", content=text) assert len(sizes) > 1, f"只调了 {len(sizes)} 次,没有分批" assert max(sizes) <= kb_mod._EMBED_BATCH, f"有一批 {max(sizes)} 条,超过了上限" async def test_progress_lands_in_the_db_while_it_runs() -> None: """进度要能被别人看见。 试过用另一个 session 单独写,但 SQLite 只允许一个写者——主事务开着的时候 那次写会被堵住,整个过程界面上都是空的、只在结束时闪一下。改成主事务 自己按批提交。 """ from app.db.models import Document seen: list[tuple[int, int]] = [] # 要切出多于一批(_EMBED_BATCH=64)的片段,否则只会有一次进度回调 text = "\n\n".join(f"第{i}段内容,这里要写得长一些才切得够多。" * 30 for i in range(300)) async with SessionLocal() as session: doc = await kb.create_document( session, collection="prog_test", title="进度", content=text, status="processing") await kb.process_document(session, doc, on_progress=lambda d, t: seen.append((d, t))) assert len(seen) > 1, "一次就跑完了,测不到进度" assert seen[-1][0] == seen[-1][1], "最后一次进度不是满的" # 关键:中途的进度确实落了库,不是只在内存里 async with SessionLocal() as session: row = await session.get(Document, doc.id) assert (row.meta or {}).get("progress", {}).get("total") == seen[-1][1] assert row.status == "ready" async def test_a_failed_document_does_not_stay_processing_forever(client) -> None: """失败了要写回 failed。不写的话它永远停在"处理中",界面上一直转圈。""" from app.db.models import Document from app.api.knowledge import _process_in_background async with SessionLocal() as session: doc = await kb.create_document( session, collection="fail_test", title="会失败的", content="一些内容", status="processing") import app.memory.kb as kb_mod async def boom(*a, **kw): raise RuntimeError("向量服务挂了") orig = kb_mod.process_document kb_mod.process_document = boom try: await _process_in_background(doc.id) finally: kb_mod.process_document = orig async with SessionLocal() as session: row = await session.get(Document, doc.id) assert row.status == "failed" assert "向量服务挂了" in row.error async def test_oversized_uploads_are_rejected_before_being_read(client) -> None: """原来是 raw = await file.read() 读完再判断大小。 传个 1GB 的文件,内存在报 413 之前就已经吃掉了——那个检查拦的是"入库", 不是"占内存"。 """ from app.core.config import settings too_big = b"x" * ((settings.max_upload_mb + 1) * 1024 * 1024) r = await client.post("/api/kb/upload", params={"collection": "big"}, files={"file": ("big.txt", too_big, "text/plain")}) assert r.status_code == 413 assert str(settings.max_upload_mb) in r.json()["detail"]