mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 16:32:00 +00:00
ThearchyHelios
62aac187d8
MCP / 自定义工具的审批三档（tools/trust.py） - 平台看不出这两类工具做什么，以前审批关卡一律当它们安全、照跑不误。现在默认每次调用都等审批； 工具页可以逐个改成「始终允许 · 门控把关」或「始终允许」，审批卡上的「始终允许」等于前者 - 门控把关：每次调用先由门控模型审一遍工具说明、参数和节点任务；判可疑、出错、超时、答不成格式， 一律交还给人，不会变成放行。判定包进 checkpoint，审批恢复重放时不再问一遍 - agent 节点、工具节点、协作团队三处关卡都接上；团队成员停不下来等人，门控没放行的调用不执行 - 每次运行在发起时快照档位（runs.tool_trust），恢复和接着跑沿用同一份；正式运行不看档位，一律等审批； 升级前发起的运行照旧不问——它们可能正停在某个审批上，多问一次答复就会对错号 - 接口：GET /api/tools 带 trust / trust_key，PUT /api/tools/trust；设置新增门控模型；新事件 tool.gated - 界面：工具页三选一、审批卡「始终允许」、设置页门控模型、运行记录解码 tool.gated； 顺手修了 MCP 工具详情 404、工具审批被驳回却显示成放行 智能体护栏（engine/guards.py），取代固定的「最多 12 步」 - 步数只留兜底：默认 100，设置里可改；以前画布自动写进节点的 12 按「没填」处理 - 同一个工具、同样的参数第二次不再执行，把上次的结果交还给它；连续 3 步拿不到新信息就收尾 - 每个节点的令牌预算（默认 200 万）和金额预算（默认不限），都可以在设置里改成不限 - 上下文用到窗口的 60% 先压缩早期的工具结果，到 85% 就收尾；临近任何一个上限先提醒模型一次 - 下游结构化校验失败时写明是上游哪个 agent 提前收了尾；复核信号带上收尾原因， 问数据页只在真是步数用满时才给「放宽步数重跑」 - 上限在发起时快照（runs.agent_limits）；升级前发起的运行一字不改走旧逻辑 检查脚本 - check-all 默认 4 道并行：量时间最多的两项互不重叠；并行时没过的最后单独重跑一遍，重跑才过的 单列成「偶发」。16 项全量从约 21 分钟降到约 7 分钟（加上一次重跑约 10 分钟）；--lanes 1 回到串行 顺带修的 - 设置接口里的 limits 改为只读：库里存着旧版设置页写进去的一份（max_agent_steps: 25），从不生效， 却盖住了真实上限，问数据页照着错的数许诺 - test_kb 里嵌入模型那条测试不再真连 api.openai.com（代理慢时一条要等近 4 分钟） - check-stream 读滚动容器前先等它挂上，并行跑时不再偶发 验证（沙箱）：tsc 干净，生产构建通过，pytest 1655 passed（两遍），check-all 16 项检查、3562 项全过； 新增后端测试 53 条，关键逻辑做了 22 处变异验证、全部抓到。
785 lines
31 KiB
Python
785 lines
31 KiB
Python
"""知识库：切块、混合检索、换 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 = """<html><head><style>p{color:red}</style></head>
|
||
<body><nav>导航</nav><h1>权限说明</h1><p>只有管理员能改角色。</p>
|
||
<script>alert(1)</script><footer>版权</footer></body></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"]
|