mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 16:52:00 +00:00
ThearchyHelios
df935599f2
三件事，都是"这三个模块一条测试都没有"的后果。 **换 embedder 会让检索整体 500。** 库里是 512 维的本地向量、查询是 1536 维的 远端向量时，cosine 直接 `shapes (1536,) and (512,) not aligned`。from_blob 有 dim 参数，但 kb.py 和 memory/store.py 调用时都没传；而且存量向量属于哪个模型 代码里没有任何地方记着，所以也无从降级。 Chunk / MemoryItem 各加 embed_model + embed_dim（走现成的 _COLUMN_MIGRATIONS， 存量回填空值，正好表示"不知道是谁建的"）。检索遇到对不上的向量整体退回纯 关键词并说破，不是把它们当 0 分静默沉底——那样用户看到的是结果莫名其妙变差， 而不是"这里有一批索引该重建了"。配套 kb.reindex + POST /api/kb/reindex。 长期记忆那份一起修了。判重那条更危险：维度碰巧相同但模型不同时 cosine 算出来 是随机数，会把不相关的旧记忆认成"重复"，然后把新内容丢掉。 **默认 embedder 没有语义能力，而开关藏在一个环境变量里。** LocalEmbedder 是 特征哈希 + TF，和 BM25 吃同一批词；alpha 默认 0.5 等于给了它一半权重，同义 改写一条都召不回。开启真模型的唯一开关是 AGENTLAB_USE_OPENAI_EMBEDDINGS—— 没界面、没人会发现。改成设置项（照 COPILOT_SETTING_KEY 的路子），知识库页顶部 能选能看能重建。 配置失败必须响：第一版写成静默退回本地，但选了 OpenAI 却拿到哈希向量的人， 在搜不出同义表达时完全无从怀疑到这里。现在抛 EmbedderUnavailable，API 400 原样交回"检查 OPENAI_API_KEY 是否配了"，且不会把配置改掉。 **检索结果不进工件库。** db_query、工具节点、agent 工具调用都落 put_json 快照， run_retrieve 什么都不落——于是结论只要来自知识库，"这个数是哪来的"就断在这儿， 而数字回指正是这个项目的签名能力。现在落 kind="retrieval_snapshot"，并发一条 retrieve.end 带 artifact id。解码层补了正式映射（原来会掉进"未知事件"兜底， 画成一坨原始 JSON）。 新增 test_kb.py 11 条，补上这三个模块此前的零覆盖，包括 embeddings.py 里 注释记着、但一直没有测试守的那个归一化边界（单条候选被 min_score 误杀）。 端到端在真实服务上验过：撞上存量向量返回 200 + 降级说明（以前 500）、 重建后 stale 归零、配不起来的远端模型 400 且配置不变。
234 lines
7.6 KiB
Python
234 lines
7.6 KiB
Python
from __future__ import annotations
|
||
|
||
import re
|
||
from typing import Any
|
||
|
||
from sqlalchemy import func, select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from app.db.models import Chunk, Document
|
||
from app.memory.embeddings import (
|
||
embed_text, embed_texts, embedder_dim, embedder_id, from_blob, hybrid_rank, to_blob, usable,
|
||
)
|
||
|
||
_TARGET = 800
|
||
_OVERLAP = 120
|
||
|
||
|
||
def chunk_text(text: str, target: int = _TARGET, overlap: int = _OVERLAP) -> list[str]:
|
||
"""按语义边界切块。
|
||
|
||
先按空行分段，段落塞不下再按句子切，实在超长才硬切。
|
||
这样切出来的块不会把一句话拦腰截断，检索命中后读起来是完整的。
|
||
"""
|
||
text = text.replace("\r\n", "\n").strip()
|
||
if not text:
|
||
return []
|
||
|
||
paragraphs = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
|
||
blocks: list[str] = []
|
||
for para in paragraphs:
|
||
if len(para) <= target:
|
||
blocks.append(para)
|
||
continue
|
||
sentences = re.split(r"(?<=[。！？.!?\n])\s*", para)
|
||
current = ""
|
||
for sentence in sentences:
|
||
if not sentence:
|
||
continue
|
||
if len(current) + len(sentence) <= target:
|
||
current += sentence
|
||
else:
|
||
if current:
|
||
blocks.append(current)
|
||
while len(sentence) > target:
|
||
blocks.append(sentence[:target])
|
||
sentence = sentence[target:]
|
||
current = sentence
|
||
if current:
|
||
blocks.append(current)
|
||
|
||
# 把相邻小块合并到接近目标长度，并留一点重叠保住上下文
|
||
merged: list[str] = []
|
||
buf = ""
|
||
for block in blocks:
|
||
if len(buf) + len(block) + 1 <= target:
|
||
buf = f"{buf}\n{block}" if buf else block
|
||
else:
|
||
if buf:
|
||
merged.append(buf)
|
||
tail = buf[-overlap:] if overlap else ""
|
||
buf = f"{tail}\n{block}" if tail else block
|
||
else:
|
||
merged.append(block)
|
||
buf = ""
|
||
if buf:
|
||
merged.append(buf)
|
||
return [m.strip() for m in merged if m.strip()]
|
||
|
||
|
||
async def ingest_document(
|
||
session: AsyncSession,
|
||
*,
|
||
collection: str,
|
||
title: str,
|
||
content: str,
|
||
source: str = "",
|
||
mime: str = "text/plain",
|
||
meta: dict[str, Any] | None = None,
|
||
) -> Document:
|
||
doc = Document(
|
||
collection=collection,
|
||
title=title or source or "未命名文档",
|
||
source=source,
|
||
mime=mime,
|
||
content=content,
|
||
meta=meta or {},
|
||
)
|
||
session.add(doc)
|
||
await session.flush()
|
||
|
||
pieces = chunk_text(content)
|
||
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},
|
||
)
|
||
)
|
||
doc.chunk_count = len(pieces)
|
||
await session.commit()
|
||
await session.refresh(doc)
|
||
return doc
|
||
|
||
|
||
async def search(
|
||
session: AsyncSession,
|
||
*,
|
||
collection: str | None,
|
||
query: str,
|
||
limit: int = 5,
|
||
alpha: float = 0.5,
|
||
on_degrade: Any = None,
|
||
) -> list[dict[str, Any]]:
|
||
"""混合检索。
|
||
|
||
存量向量和当前 embedder 对不上时**整个查询退回纯关键词**，并通过
|
||
on_degrade 把原因交出去。不把对不上的那些当 0 分：那会让它们静默沉底，
|
||
用户看到的是结果莫名其妙变差，而不是"这里有一批索引该重建了"。
|
||
|
||
在此之前这里连维度都不查，换个 embedder 直接 ValueError
|
||
（shapes (1536,) and (512,) not aligned），整个检索 500。
|
||
"""
|
||
stmt = select(Chunk)
|
||
if collection:
|
||
stmt = stmt.where(Chunk.collection == collection)
|
||
rows = list((await session.execute(stmt)).scalars())
|
||
if not rows:
|
||
return []
|
||
|
||
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})
|
||
on_degrade(
|
||
f"{collection or '全部集合'}里有 {len(stale)}/{len(rows)} 条向量由"
|
||
f"「{'、'.join(names)}」建立，和当前的「{embedder_id()}」对不上，"
|
||
f"本次已退回纯关键词检索。重建索引后恢复。"
|
||
)
|
||
alpha = 0.0 # 纯关键词
|
||
vectors: list[Any] = [None] * len(rows)
|
||
query_vec = await embed_text(query)
|
||
else:
|
||
query_vec = await embed_text(query)
|
||
vectors = [from_blob(r.embedding, embedder_dim()) for r in rows]
|
||
|
||
ranked = hybrid_rank(
|
||
query,
|
||
[r.content for r in rows],
|
||
vectors,
|
||
query_vec,
|
||
alpha=alpha,
|
||
)
|
||
out: list[dict[str, Any]] = []
|
||
for idx, score, parts in ranked[:limit]:
|
||
row = rows[idx]
|
||
out.append(
|
||
{
|
||
"chunk_id": row.id,
|
||
"document_id": row.document_id,
|
||
"title": (row.meta or {}).get("title", ""),
|
||
"ordinal": row.ordinal,
|
||
"content": row.content,
|
||
"score": round(score, 4),
|
||
"signals": {k: round(v, 4) for k, v in parts.items()},
|
||
}
|
||
)
|
||
return out
|
||
|
||
|
||
async def reindex(session: AsyncSession, collection: str | None = None) -> dict[str, Any]:
|
||
"""用当前 embedder 重算向量。
|
||
|
||
换了 embedding 模型之后唯一的恢复手段。按批做：一次几千条 chunk 全塞给
|
||
远端 embedding 接口，要么超时要么被限流，而中途失败又没有断点。
|
||
"""
|
||
stmt = select(Chunk)
|
||
if collection:
|
||
stmt = stmt.where(Chunk.collection == collection)
|
||
rows = list((await session.execute(stmt)).scalars())
|
||
if not rows:
|
||
return {"collection": collection, "reindexed": 0, "embedder": embedder_id()}
|
||
|
||
model_id, dim = embedder_id(), embedder_dim()
|
||
batch = 64
|
||
for i in range(0, len(rows), batch):
|
||
part = rows[i:i + batch]
|
||
vectors = await embed_texts([r.content for r in part])
|
||
for row, vec in zip(part, vectors):
|
||
row.embedding = to_blob(vec)
|
||
row.embed_model = model_id
|
||
row.embed_dim = dim
|
||
await session.commit()
|
||
return {"collection": collection, "reindexed": len(rows), "embedder": model_id}
|
||
|
||
|
||
async def stale_count(session: AsyncSession, collection: str | None = None) -> int:
|
||
"""有多少条向量和当前 embedder 对不上。知识库页据此提示要不要重建。"""
|
||
stmt = select(Chunk)
|
||
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))
|
||
|
||
|
||
async def delete_document(session: AsyncSession, doc_id: str) -> bool:
|
||
doc = await session.get(Document, doc_id)
|
||
if not doc:
|
||
return False
|
||
await session.delete(doc)
|
||
await session.commit()
|
||
return True
|
||
|
||
|
||
async def list_collections(session: AsyncSession) -> list[dict[str, Any]]:
|
||
rows = await session.execute(
|
||
select(
|
||
Document.collection,
|
||
func.count(Document.id),
|
||
func.sum(Document.chunk_count),
|
||
).group_by(Document.collection)
|
||
)
|
||
return [
|
||
{"collection": name, "documents": docs, "chunks": chunks or 0}
|
||
for name, docs, chunks in rows
|
||
]
|