dadao
d079564d6c
CI / backend (push) Canceled after 0s
之前知识只能"存"，不能"找"。这一版把链路补上：切片 → 向量 → 打分 → 带来源返回。
- chunk 表 + 迁移 0006：正文按标题归段切、代码块整块保留、超长按句切、
相邻片重叠 80 字；RLS 与 GRANT 与其它租户表一致
- 向量提供方做成协议：内置 hash（字符二元组，零依赖可复现，作为基线）
与 OpenAI 兼容端点两种，换实现不动业务代码
- 检索是混合打分：向量 0.6 + pg_trgm 0.4。专有编号、单号这类必须原样命中的词
靠词法兜，换个说法表达的同一件事靠向量兜
- 三条纪律：只搜 published（过滤写在 SQL 里）、低于阈值宁可不返回、
每条命中都带来源标题与 item_id
- 发布时自动建索引，改完正文用 POST /knowledge/{id}/reindex 重建
- Web 知识页加检索框与命中卡片；新增 e2e/retrieval.mjs 6 项
为什么不是 pgvector：本机 Docker 起不来、fnsvr-01 没有该扩展且明确不动。
先用 real[] + SQL 余弦把链路跑通，将来换列类型与索引即可，接口不变。
后端 79 passed，e2e smoke/knowledge/richtext/retrieval 全绿。
209 lines
6.9 KiB
Python
209 lines
6.9 KiB
Python
"""知识的切片索引与检索（M3）。
|
||
|
||
两条硬规则：
|
||
|
||
1. **只有 `published` 的知识能被检索到**。草稿是"还没被人认可的话"，
|
||
让它进检索等于把猜测当结论引用出去。过滤写在 SQL 里，不靠调用方自觉。
|
||
2. **检索必须带来源**。每个命中都返回 `item_id` 与标题，用户能点回去看原文；
|
||
这和 D3（知识必须有来源）是同一条原则的延伸——不给"出处不存在"的答案。
|
||
|
||
评分是向量与词法的加权和：
|
||
- 向量负责"说的是同一件事但换了词"；
|
||
- 词法（pg_trgm 的 trigram 相似度）负责"专有名词、单号、人名"这类必须原样命中的词，
|
||
而哈希向量对这类词并不敏感。
|
||
两者都不足时宁可返回空，也不硬凑一条不相关的结果——检索失败比误导便宜。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import uuid
|
||
from typing import Any
|
||
|
||
from sqlalchemy import delete, select, text
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from ..config import get_settings
|
||
from ..domain.chunking import split_markdown
|
||
from ..errors import AppError, ErrorCode
|
||
from ..models_knowledge import Chunk, KnowledgeDoc
|
||
from ..models_task import Item
|
||
from .embedding import get_provider
|
||
|
||
SEARCH_SQL = text(
|
||
"""
|
||
SELECT c.item_id,
|
||
i.title,
|
||
k.doc_type,
|
||
c.idx,
|
||
c.heading,
|
||
c.content,
|
||
coalesce(app_cosine(c.embedding, CAST(:qvec AS real[])), 0) AS vec_score,
|
||
coalesce(similarity(c.content, :q), 0) AS lex_score
|
||
FROM chunk c
|
||
JOIN item i ON i.id = c.item_id AND i.workspace_id = c.workspace_id
|
||
JOIN knowledge_doc k ON k.item_id = c.item_id AND k.workspace_id = c.workspace_id
|
||
WHERE c.workspace_id = :ws
|
||
AND k.status_hint = 'published'
|
||
ORDER BY (
|
||
coalesce(app_cosine(c.embedding, CAST(:qvec AS real[])), 0) * :vw
|
||
+ coalesce(similarity(c.content, :q), 0) * :lw
|
||
) DESC, c.item_id, c.idx
|
||
LIMIT :limit
|
||
"""
|
||
)
|
||
|
||
|
||
def _vector_param(values: list[float]) -> list[float]:
|
||
"""把向量交给驱动去编码。
|
||
|
||
一开始我拼成 `'{0.1,0.2}'` 字符串再 `CAST(... AS real[])`，asyncpg 会因为 CAST
|
||
已经告诉它"这是数组"而拒绝字符串（`a sized iterable container expected`）。
|
||
直接传 list 最省事；非有限值替换成 0，免得 NaN 进 SQL 变成错误输入。
|
||
"""
|
||
return [value if value == value and abs(value) != float("inf") else 0.0 for value in values]
|
||
|
||
|
||
async def reindex_document(
|
||
session: AsyncSession,
|
||
ws_id: uuid.UUID,
|
||
item_id: uuid.UUID,
|
||
) -> int:
|
||
"""重建一份知识的切片，返回切片数。
|
||
|
||
全量重建而不是增量：知识正文改动通常牵动上下几段，增量对齐的复杂度
|
||
换不来多少收益（一份知识撑死几十片）。
|
||
"""
|
||
item = (
|
||
await session.execute(select(Item).where(Item.workspace_id == ws_id, Item.id == item_id))
|
||
).scalar_one_or_none()
|
||
if item is None:
|
||
raise AppError(ErrorCode.ITEM_NOT_FOUND, "这件事不存在")
|
||
|
||
doc = (
|
||
await session.execute(
|
||
select(KnowledgeDoc).where(
|
||
KnowledgeDoc.workspace_id == ws_id,
|
||
KnowledgeDoc.item_id == item_id,
|
||
)
|
||
)
|
||
).scalar_one_or_none()
|
||
if doc is None:
|
||
raise AppError(ErrorCode.KNOWLEDGE_NOT_FOUND, "这份知识不存在")
|
||
|
||
settings = get_settings()
|
||
source = "\n\n".join(part for part in (item.title, item.body_md or "") if part)
|
||
drafts = split_markdown(
|
||
source,
|
||
max_chars=settings.chunk_max_chars,
|
||
overlap=settings.chunk_overlap_chars,
|
||
)
|
||
|
||
provider = get_provider(settings)
|
||
vectors = await provider.embed([draft.content for draft in drafts]) if drafts else []
|
||
|
||
await session.execute(
|
||
delete(Chunk).where(Chunk.workspace_id == ws_id, Chunk.item_id == item_id)
|
||
)
|
||
for index, (draft, vector) in enumerate(zip(drafts, vectors, strict=True)):
|
||
session.add(
|
||
Chunk(
|
||
workspace_id=ws_id,
|
||
item_id=item_id,
|
||
idx=index,
|
||
heading=draft.heading,
|
||
content=draft.content,
|
||
char_count=len(draft.content),
|
||
embedding=vector,
|
||
embed_model=provider.name,
|
||
)
|
||
)
|
||
doc.chunk_count = len(drafts)
|
||
await session.flush()
|
||
return len(drafts)
|
||
|
||
|
||
async def search_knowledge(
|
||
session: AsyncSession,
|
||
ws_id: uuid.UUID,
|
||
query: str,
|
||
*,
|
||
limit: int = 10,
|
||
mode: str = "hybrid",
|
||
) -> list[dict[str, Any]]:
|
||
"""检索已发布知识，返回按相关度排序的切片。
|
||
|
||
`mode`: hybrid（默认）/ vector / keyword —— 前两者主要用于对比效果。
|
||
"""
|
||
cleaned = (query or "").strip()
|
||
if not cleaned:
|
||
return []
|
||
|
||
settings = get_settings()
|
||
provider = get_provider(settings)
|
||
query_vector = (await provider.embed([cleaned]))[0]
|
||
|
||
if mode == "vector":
|
||
vector_weight, lexical_weight = 1.0, 0.0
|
||
elif mode == "keyword":
|
||
vector_weight, lexical_weight = 0.0, 1.0
|
||
else:
|
||
vector_weight, lexical_weight = (
|
||
settings.search_vector_weight,
|
||
settings.search_lexical_weight,
|
||
)
|
||
|
||
rows = (
|
||
await session.execute(
|
||
SEARCH_SQL,
|
||
{
|
||
"qvec": _vector_param(query_vector),
|
||
"q": cleaned,
|
||
"ws": str(ws_id),
|
||
"vw": vector_weight,
|
||
"lw": lexical_weight,
|
||
"limit": limit,
|
||
},
|
||
)
|
||
).mappings()
|
||
|
||
hits: list[dict[str, Any]] = []
|
||
for row in rows:
|
||
vector_score = float(row["vec_score"] or 0.0)
|
||
lexical_score = float(row["lex_score"] or 0.0)
|
||
hits.append(
|
||
{
|
||
"item_id": row["item_id"],
|
||
"title": row["title"],
|
||
"doc_type": row["doc_type"],
|
||
"idx": row["idx"],
|
||
"heading": row["heading"],
|
||
"content": row["content"],
|
||
"score": round(vector_score * vector_weight + lexical_score * lexical_weight, 4),
|
||
"vector_score": round(vector_score, 4),
|
||
"lexical_score": round(lexical_score, 4),
|
||
}
|
||
)
|
||
hits = [hit for hit in hits if hit["score"] >= settings.search_min_score]
|
||
return hits
|
||
|
||
|
||
async def reindex_workspace(session: AsyncSession, ws_id: uuid.UUID) -> int:
|
||
"""把一个工作区里已发布的知识全部重建，返回处理的知识份数。
|
||
|
||
换 embedding 模型之后必须整体跑一次，否则新旧向量混在一张表里，
|
||
相似度分数没法比较。
|
||
"""
|
||
item_ids = (
|
||
await session.execute(
|
||
select(KnowledgeDoc.item_id).where(
|
||
KnowledgeDoc.workspace_id == ws_id,
|
||
KnowledgeDoc.status_hint == "published",
|
||
)
|
||
)
|
||
).scalars()
|
||
count = 0
|
||
for item_id in item_ids:
|
||
await reindex_document(session, ws_id, item_id)
|
||
count += 1
|
||
return count
|