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]
|