76 lines
3.2 KiB
Python
76 lines
3.2 KiB
Python
"""Embedding 服务（T21-1 · FLOW §3「embedding_tool(Ollama bge-m3)」）。
|
||
|
||
技术选型（01-技术栈与版本.md §4）：Ollama + bge-m3，**1024 维**，文档向量
|
||
不出内网。调用 Ollama 批量接口 ``POST /api/embed``（一次请求携带多段文本，
|
||
比逐条 /api/embeddings 少一个数量级的往返开销）。
|
||
|
||
失败口径（拍板 2026-09-07）：Ollama 连不上 / 模型缺失 / 响应结构异常 /
|
||
维度不符一律抛 EmbeddingError——**禁止静默返回零向量或截断降级**。向量
|
||
错误会静默污染检索结果（查得慢、查得偏都难归因），明确失败比带病成功
|
||
更可运维；调用方（build_kb 入库 / rag_service 检索）按明确错误处理。
|
||
|
||
测试：单测经 httpx.MockTransport 模拟 Ollama 响应，不依赖 Ollama 进程；
|
||
真库联调归 scripts/kb/build_kb.py 与收尾验证。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import httpx
|
||
|
||
from app.config.settings import settings
|
||
|
||
|
||
class EmbeddingError(Exception):
|
||
"""向量化失败（连接/模型/响应/维度）——调用方明确处理，不降级假数据。"""
|
||
|
||
|
||
def _client() -> httpx.Client:
|
||
"""HTTP 客户端工厂（测试 monkeypatch 点；超时取 settings）。"""
|
||
return httpx.Client(timeout=settings.embed_timeout_seconds)
|
||
|
||
|
||
def embed_texts(texts: list[str]) -> list[list[float]]:
|
||
"""批量文本 → 向量（顺序与入参一一对应）。
|
||
|
||
空列表直接返回 []（不发请求，防 Ollama 对空 input 报错）；单条空串
|
||
交由 Ollama 侧语义处理（bge-m3 对空串也能出向量，调用方自行决定是否
|
||
过滤——切块层保证不为空，此处不做第二套校验）。
|
||
"""
|
||
if not texts:
|
||
return []
|
||
payload = {"model": settings.embed_model, "input": list(texts)}
|
||
try:
|
||
with _client() as client:
|
||
resp = client.post(f"{settings.ollama_base_url}/api/embed", json=payload)
|
||
except httpx.HTTPError as exc:
|
||
raise EmbeddingError(
|
||
f"Ollama 连接失败（{settings.ollama_base_url}）：{exc}"
|
||
) from exc
|
||
if resp.status_code != 200:
|
||
# 模型未拉取时 Ollama 返回 404 + {"error": "model '...' not found"}
|
||
detail = (resp.text or "")[:200]
|
||
raise EmbeddingError(
|
||
f"Ollama embedding 请求失败 HTTP {resp.status_code}（模型 "
|
||
f"{settings.embed_model} 是否已 pull？）：{detail}"
|
||
)
|
||
body = resp.json()
|
||
embeddings = (body or {}).get("embeddings")
|
||
if not isinstance(embeddings, list) or len(embeddings) != len(texts):
|
||
raise EmbeddingError(
|
||
f"Ollama 响应结构异常：embeddings 数量与输入不匹配"
|
||
f"（输入 {len(texts)} 条）"
|
||
)
|
||
for vec in embeddings:
|
||
if not isinstance(vec, list) or len(vec) != settings.embed_dim:
|
||
raise EmbeddingError(
|
||
f"向量维度不符：期望 {settings.embed_dim}，实际 "
|
||
f"{len(vec) if isinstance(vec, list) else '非数组'}"
|
||
f"——检查 embed_dim 配置与 {settings.embed_model} 输出是否一致"
|
||
)
|
||
return embeddings
|
||
|
||
|
||
def embed_text(text: str) -> list[float]:
|
||
"""单文本便捷方法（= embed_texts([text])[0]）。"""
|
||
return embed_texts([text])[0]
|