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 全绿。
132 lines
4.7 KiB
Python
132 lines
4.7 KiB
Python
"""文本向量化（M3）。
|
||
|
||
为什么要抽象一层：今天用内置的哈希向量，明天换成 OpenAI 兼容的 `/embeddings`，
|
||
后天可能换成自建模型。检索代码只依赖 `embed() -> list[float]`，换实现不动业务。
|
||
|
||
内置 `HashEmbedding` 的诚实说明：它不是"语义模型"，是**字符 n-gram 的哈希词袋**。
|
||
它能把同词、近义写法（"发票"/"票据"这种）之外的东西区分开，抓不到真正的语义相似。
|
||
之所以先上它：本机没有可用的向量库与模型服务，检索链路不能一直不落地；
|
||
而且它零依赖、可复现（同样的文本永远同样的向量），适合做测试基线。
|
||
真上线前应当把 `EMBEDDING_PROVIDER` 换成 `openai`，并把已发布知识整体重建一次。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import math
|
||
import re
|
||
from collections.abc import Sequence
|
||
from typing import Protocol
|
||
|
||
import httpx
|
||
|
||
from ..config import Settings, get_settings
|
||
|
||
CJK_RE = re.compile(r"[\u4e00-\u9fff]")
|
||
WORD_RE = re.compile(r"[a-z0-9]+")
|
||
|
||
|
||
class EmbeddingProvider(Protocol):
|
||
"""向量提供方。`name` 会写进切片，换模型后能看出哪些是旧向量。"""
|
||
|
||
name: str
|
||
|
||
async def embed(self, texts: Sequence[str]) -> list[list[float]]: ...
|
||
|
||
|
||
def _tokens(text: str) -> list[str]:
|
||
"""中文按字二元组、英文按词。
|
||
|
||
中文不用分词器：单字太碎（"发""票"都对不上"发票"），二元组能拿到大多数
|
||
有意义的搭配，又不需要维护词典。
|
||
"""
|
||
lowered = text.lower()
|
||
tokens = WORD_RE.findall(lowered)
|
||
chars = [c for c in lowered if CJK_RE.match(c)]
|
||
tokens.extend(chars[i] + chars[i + 1] for i in range(len(chars) - 1))
|
||
tokens.extend(chars)
|
||
return tokens
|
||
|
||
|
||
class HashEmbedding:
|
||
"""字符 n-gram 的带符号哈希词袋，L2 归一化。"""
|
||
|
||
def __init__(self, dim: int = 256) -> None:
|
||
self.dim = dim
|
||
self.name = f"hash-{dim}"
|
||
|
||
async def embed(self, texts: Sequence[str]) -> list[list[float]]:
|
||
return [self._one(text) for text in texts]
|
||
|
||
def _one(self, text: str) -> list[float]:
|
||
vector = [0.0] * self.dim
|
||
counts: dict[str, int] = {}
|
||
for token in _tokens(text):
|
||
counts[token] = counts.get(token, 0) + 1
|
||
for token, count in counts.items():
|
||
digest = hashlib.blake2b(token.encode("utf-8"), digest_size=8).digest()
|
||
raw = int.from_bytes(digest, "big")
|
||
index = raw % self.dim
|
||
# 用最高位决定正负：无符号哈希会让不相干的词在向量里越加越像
|
||
sign = 1.0 if (raw >> 63) & 1 else -1.0
|
||
vector[index] += sign * (1.0 + math.log(count))
|
||
norm = math.sqrt(sum(v * v for v in vector))
|
||
if norm == 0:
|
||
return vector
|
||
return [v / norm for v in vector]
|
||
|
||
|
||
class OpenAIEmbedding:
|
||
"""任何 OpenAI 兼容的 `/embeddings` 端点（BYOK）。"""
|
||
|
||
def __init__(self, *, base_url: str, api_key: str, model: str) -> None:
|
||
self.base_url = base_url.rstrip("/")
|
||
self.api_key = api_key
|
||
self.model = model
|
||
self.name = f"openai:{model}"
|
||
|
||
async def embed(self, texts: Sequence[str]) -> list[list[float]]:
|
||
if not texts:
|
||
return []
|
||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||
response = await client.post(
|
||
f"{self.base_url}/embeddings",
|
||
headers={"Authorization": f"Bearer {self.api_key}"},
|
||
json={"model": self.model, "input": list(texts)},
|
||
)
|
||
response.raise_for_status()
|
||
payload = response.json()
|
||
# 按 index 排序：并发端点的返回顺序没有保证
|
||
rows = sorted(payload["data"], key=lambda row: row["index"])
|
||
return [row["embedding"] for row in rows]
|
||
|
||
|
||
_provider: EmbeddingProvider | None = None
|
||
|
||
|
||
def get_provider(settings: Settings | None = None) -> EmbeddingProvider:
|
||
"""进程内单例：哈希实现无状态，OpenAI 实现也不需要每次新建连接池。"""
|
||
global _provider
|
||
if _provider is not None:
|
||
return _provider
|
||
settings = settings or get_settings()
|
||
if settings.embedding_provider == "openai":
|
||
if not settings.embedding_api_key:
|
||
raise RuntimeError(
|
||
"EMBEDDING_PROVIDER=openai 但没有配置 EMBEDDING_API_KEY；"
|
||
"要么补上密钥，要么改回 hash"
|
||
)
|
||
_provider = OpenAIEmbedding(
|
||
base_url=settings.embedding_base_url,
|
||
api_key=settings.embedding_api_key,
|
||
model=settings.embedding_model,
|
||
)
|
||
else:
|
||
_provider = HashEmbedding()
|
||
return _provider
|
||
|
||
|
||
def reset_provider_for_tests() -> None:
|
||
global _provider
|
||
_provider = None
|