jerryandClaude Fable 5
5db3ead73e
- vstore 独立安装包(8094):解析/嵌入/加密 vault/faiss 检索,存储路径支持本地 file 与 http(s):// URL,启动握手校验,集群化预留(embed 独立模块) - BFF docgen 路由:上传入库/generate/编辑/导出/下载/主库管理,意图护栏(与前端同词典),vstore 失败兜底,quota 507 与 TTL 惰性清理 - 前端:自然触发意图路由、拖拽上传 chips、能力徽标、右侧文档 TAB(WYSIWYG 编辑/prompt 修改/PDF+DOCX 导出下载)、右栏拖宽+手柄仿左栏 - 批量导入工具 docgen_import.py(并发 8、导入即加密+删源明文) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
92 lines
3.1 KiB
Python
92 lines
3.1 KiB
Python
"""embed.py —— ollama 嵌入调用封装（httpx，并发 ≤8，bge-m3 1024 维）。
|
||
|
||
- embed_batch(texts)：同步函数，分批并发调 POST {url}/api/embed
|
||
- 空文本返回零向量；维度不符抛 EmbedError
|
||
- 单元测试注入假 embed（monkeypatch），本模块不依赖真 ollama
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
|
||
import httpx
|
||
|
||
from .settings import (
|
||
DEFAULT_EMBED_MODEL,
|
||
DEFAULT_EMBED_URL,
|
||
EMBED_CONCURRENCY,
|
||
EMBED_DIM,
|
||
EMBED_TIMEOUT,
|
||
get_paths,
|
||
)
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class EmbedError(Exception):
|
||
"""嵌入调用失败（服务不可达/维度不符/响应异常）。"""
|
||
|
||
|
||
def _zero_vector() -> list[float]:
|
||
return [0.0] * EMBED_DIM
|
||
|
||
|
||
def _request_one(client: httpx.Client, url: str, model: str, text: str) -> list[float]:
|
||
"""单条嵌入请求。空文本返回零向量（不打上游）。"""
|
||
if not text or not text.strip():
|
||
return _zero_vector()
|
||
try:
|
||
resp = client.post(
|
||
f"{url.rstrip('/')}/api/embed",
|
||
json={"model": model, "input": text},
|
||
timeout=EMBED_TIMEOUT,
|
||
)
|
||
resp.raise_for_status()
|
||
payload = resp.json()
|
||
embeddings = payload.get("embeddings")
|
||
if not embeddings or not isinstance(embeddings[0], list):
|
||
raise EmbedError(f"嵌入响应形态异常：{str(payload)[:200]}")
|
||
vector = [float(v) for v in embeddings[0]]
|
||
if len(vector) != EMBED_DIM:
|
||
raise EmbedError(f"嵌入维度不符：期望 {EMBED_DIM}，实际 {len(vector)}")
|
||
return vector
|
||
except EmbedError:
|
||
raise
|
||
except (httpx.HTTPError, ValueError) as exc:
|
||
raise EmbedError(f"嵌入服务调用失败：{exc}") from exc
|
||
|
||
|
||
def embed_batch(
|
||
texts: list[str],
|
||
url: str | None = None,
|
||
model: str | None = None,
|
||
) -> list[list[float]]:
|
||
"""批量嵌入：线程池并发 ≤ EMBED_CONCURRENCY，保持输入顺序。
|
||
|
||
url/model 缺省取 settings（VSTORE_EMBED_URL / VSTORE_EMBED_MODEL）。
|
||
"""
|
||
if not texts:
|
||
return []
|
||
paths = get_paths()
|
||
base_url = url or paths.embed_url or DEFAULT_EMBED_URL
|
||
model_name = model or paths.embed_model or DEFAULT_EMBED_MODEL
|
||
|
||
results: list[list[float] | None] = [None] * len(texts)
|
||
with httpx.Client() as client:
|
||
with ThreadPoolExecutor(max_workers=EMBED_CONCURRENCY) as pool:
|
||
futures = {
|
||
pool.submit(_request_one, client, base_url, model_name, text): idx
|
||
for idx, text in enumerate(texts)
|
||
}
|
||
for future in futures:
|
||
idx = futures[future]
|
||
try:
|
||
results[idx] = future.result()
|
||
except EmbedError as exc:
|
||
logger.warning("embed: 第 %d 条嵌入失败：%s", idx, exc)
|
||
results[idx] = None
|
||
except Exception as exc: # 线程池其他异常也降级为零向量
|
||
logger.warning("embed: 第 %d 条嵌入异常：%s", idx, exc)
|
||
results[idx] = None
|
||
return [v if v is not None else _zero_vector() for v in results]
|