PlantAI-rag-new/vstore/embed.py

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