"""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): """嵌入调用失败（服务不可达/维度不符/响应异常）。""" # 测试注入钩子（仿 vault.set_embed_fn 模式）：T1 用例按契约名 _CHUNKED_REQUEST_FN # monkeypatch 本模块该函数后，embed_batch_chunked 不触网。 def _CHUNKED_REQUEST_FN(payload: dict) -> list[list[float]]: """单批 /api/embed 请求（input 数组）。返回该批 embedding 列表（真 Ollama HTTP）。""" paths = get_paths() url = (payload.get("url") or paths.embed_url or DEFAULT_EMBED_URL).rstrip("/") model = payload.get("model") or paths.embed_model or DEFAULT_EMBED_MODEL inputs = [str(t or "") for t in payload.get("input") or []] try: resp = httpx.post( f"{url}/api/embed", json={"model": model, "input": inputs}, timeout=EMBED_TIMEOUT, ) resp.raise_for_status() data = resp.json() embeddings = data.get("embeddings") if not isinstance(embeddings, list): raise EmbedError(f"嵌入响应形态异常：{str(data)[:200]}") out: list[list[float]] = [] for vec in embeddings: if not isinstance(vec, list): raise EmbedError(f"嵌入响应形态异常：{str(data)[:200]}") vector = [float(v) for v in vec] if len(vector) != EMBED_DIM: raise EmbedError(f"嵌入维度不符：期望 {EMBED_DIM}，实际 {len(vector)}") out.append(vector) if len(out) != len(inputs): raise EmbedError(f"嵌入条数不符：期望 {len(inputs)}，实际 {len(out)}") return out except EmbedError: raise except (httpx.HTTPError, ValueError) as exc: raise EmbedError(f"嵌入服务调用失败：{exc}") from exc CHUNKED_BATCH_SIZE = 16 # embed_batch_chunked 每批条数（任务书契约） CHUNKED_WORKERS = 4 # 批并发线程数（任务书契约） 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] def embed_batch_chunked( texts: list[str], batch_size: int = CHUNKED_BATCH_SIZE, workers: int = CHUNKED_WORKERS, url: str | None = None, model: str | None = None, ) -> list[list[float]]: """chunk 批量嵌入：/api/embed 的 input 数组（batch 16 × workers 4 并发）。 与既有 embed_batch（单条循环、并发 8）区分。嵌入纪律（FR-PRIV-002）： 单批失败整批重试一次，仍失败该批全部降级零向量；保持输入顺序； 批间并发上限 = workers（线程数）。 """ if not texts: return [] batch_size = max(1, int(batch_size)) workers = max(1, int(workers)) batches = [texts[i : i + batch_size] for i in range(0, len(texts), batch_size)] results: list[list[float]] = [None] * len(texts) def run_batch(offset: int, batch: list[str]) -> None: payload = {"input": list(batch), "url": url, "model": model} vecs: list[list[float]] | None = None try: vecs = _CHUNKED_REQUEST_FN(payload) except Exception as exc: logger.warning("chunks: 批量嵌入失败（第 %d 批，首调）：%s", offset, exc) try: vecs = _CHUNKED_REQUEST_FN(payload) except Exception as exc2: logger.warning( "chunks: 批量嵌入重试仍失败（第 %d 批），降级零向量：%s", offset, exc2 ) vecs = None if vecs is None: for i in range(len(batch)): results[offset + i] = _zero_vector() return for i, vec in enumerate(vecs): if offset + i >= len(results): break if isinstance(vec, list) and len(vec) == EMBED_DIM: results[offset + i] = vec else: logger.warning("chunks: 第 %d 条嵌入维度不符，降级零向量", offset + i) results[offset + i] = _zero_vector() for i in range(len(vecs), len(batch)): # 响应条数不足：缺的降级零向量 if offset + i < len(results): results[offset + i] = _zero_vector() with ThreadPoolExecutor(max_workers=workers) as pool: futures = [] offset = 0 for batch in batches: futures.append(pool.submit(run_batch, offset, batch)) offset += len(batch) for future in futures: try: future.result() except Exception as exc: # 线程池兜底异常：整批降级零向量 logger.warning("chunks: 批量嵌入线程异常：%s", exc) return [v if v is not None else _zero_vector() for v in results]