""" Embedding 封装：通过 OpenAI 兼容 ``/v1/embeddings`` 远程 API 获取向量 （ModelScope / OpenAI / DashScope 等由环境变量选择）。 """ import gc import os from dataclasses import dataclass from typing import Any, List, Optional, Union import numpy as np import logging logger = logging.getLogger(__name__) @dataclass class _RemoteEmbeddingEnv: """从环境变量解析出的远程 OpenAI-Compatible Embedding 配置。""" api_key: str base_url: str model: str max_batch: int label: str def _remote_embedding_from_env() -> _RemoteEmbeddingEnv: """ 优先 ModelScope（MODELSCOPE_*）；未配置时回退 DashScope（DASHSCOPE_*）。 """ ms_key = os.getenv("MODELSCOPE_API_KEY", "").strip() ms_base = os.getenv("MODELSCOPE_BASE_URL", "").strip() ds_key = os.getenv("DASHSCOPE_API_KEY", "").strip() ds_base = ( os.getenv("DASHSCOPE_BASE_URL") or os.getenv("DASHSCOPE_base_url", "") ).strip() # 仅当配置了 API Key 时走 ModelScope（避免仅有 BASE_URL 时误判、阻断 DashScope） if ms_key: base_url = ms_base or "https://api-inference.modelscope.cn/v1" model = ( os.getenv("MODELSCOPE_EMBEDDING_MODEL") or os.getenv("MODELSCOPE_MODEL", "Qwen/Qwen3-Embedding-8B") ).strip() mb = os.getenv("MODELSCOPE_EMBEDDING_MAX_BATCH", "32").strip() max_batch = max(1, int(mb)) if mb.isdigit() else 32 if not model: raise ValueError("未配置 MODELSCOPE_EMBEDDING_MODEL（或 MODELSCOPE_MODEL）") return _RemoteEmbeddingEnv( api_key=ms_key, base_url=base_url.rstrip("/"), model=model, max_batch=max_batch, label="ModelScope", ) # OpenAI 官方或兼容网关：OPENAI_API_KEY、OPENAI_EMBEDDING_MODEL、可选 OPENAI_BASE_URL oa_key = os.getenv("OPENAI_API_KEY", "").strip() if oa_key: if oa_key.startswith("http://") or oa_key.startswith("https://"): raise ValueError( "OPENAI_API_KEY 不能填写为 URL：请将网关地址写到 OPENAI_BASE_URL" "（例如 http://host:9080/v1），密钥单独写在 OPENAI_API_KEY" ) oa_base = ( os.getenv("OPENAI_BASE_URL", "").strip() or "https://api.openai.com/v1" ) oa_model = os.getenv("OPENAI_EMBEDDING_MODEL", "").strip() if not oa_model: raise ValueError( "使用 OpenAI 兼容 Embedding 时请设置 OPENAI_EMBEDDING_MODEL" ) mb = os.getenv("OPENAI_EMBEDDING_MAX_BATCH", "100").strip() max_batch = max(1, int(mb)) if mb.isdigit() else 100 return _RemoteEmbeddingEnv( api_key=oa_key, base_url=oa_base.rstrip("/"), model=oa_model, max_batch=max_batch, label="OpenAI", ) if not ds_key: raise ValueError( "远程 Embedding 未配置：请设置 MODELSCOPE_API_KEY（及可选 BASE_URL），" "或 OPENAI_API_KEY / OPENAI_EMBEDDING_MODEL（及可选 OPENAI_BASE_URL），" "或 DASHSCOPE_API_KEY / DASHSCOPE_BASE_URL / DASHSCOPE_MODEL" ) if not ds_base: raise ValueError("未配置 DASHSCOPE_BASE_URL（或 DASHSCOPE_base_url）") model = os.getenv("DASHSCOPE_MODEL", "").strip() if not model: raise ValueError("未配置 DASHSCOPE_MODEL") mb = os.getenv("DASHSCOPE_EMBEDDING_MAX_BATCH", "10").strip() max_batch = max(1, int(mb)) if mb.isdigit() else 10 return _RemoteEmbeddingEnv( api_key=ds_key, base_url=ds_base.rstrip("/"), model=model, max_batch=max_batch, label="DashScope", ) class OpenAICompatibleRemoteEmbedding: def __init__( self, api_key: Optional[str] = None, base_url: Optional[str] = None, model: Optional[str] = None, max_batch: Optional[int] = None, provider_label: Optional[str] = None, ): try: from openai import OpenAI except ImportError as e: raise ImportError( "使用远程 Embedding 需要安装 openai：pip install openai" ) from e if api_key is not None and base_url is not None and model is not None: cfg = _RemoteEmbeddingEnv( api_key=api_key.strip(), base_url=base_url.strip().rstrip("/"), model=model.strip(), max_batch=max(1, int(max_batch)) if max_batch is not None else 32, label=provider_label or "custom", ) else: cfg = _remote_embedding_from_env() self.api_key = cfg.api_key self.base_url = cfg.base_url self.model = cfg.model self._api_max_batch = cfg.max_batch self._provider_label = cfg.label self._client = OpenAI(api_key=self.api_key, base_url=self.base_url) vd = os.getenv("VECTOR_DIM", "").strip() self._embedding_dim: Optional[int] = int(vd) if vd.isdigit() else None logger.info( "使用 %s Embedding API：model=%s，base_url=%s，max_batch=%s", self._provider_label, self.model, self.base_url, self._api_max_batch, ) @property def embedding_dim(self) -> int: if self._embedding_dim is None: raise RuntimeError( "尚未获知向量维度：请先执行一次 encode，或在 .env 中设置 VECTOR_DIM" ) return self._embedding_dim def _set_dim_from_vector(self, vec: List[float]) -> None: if self._embedding_dim is None: self._embedding_dim = len(vec) logger.info("[OK] Embedding 向量维度：%s", self._embedding_dim) def encode( self, texts: Union[str, List[str]], batch_size: int = 10, normalize: bool = True, max_length: int = 8192, show_progress: bool = False, ) -> np.ndarray: del max_length # API 侧截断，此处仅保持签名与历史调用方一致 if isinstance(texts, str): texts = [texts] if not texts: dim = self._embedding_dim if dim is None: vd = os.getenv("VECTOR_DIM", "").strip() dim = int(vd) if vd.isdigit() else 1024 return np.empty((0, dim), dtype=np.float32) # 无论调用方传多大，不能超过远端接口单次条数上限 step = max(1, min(int(batch_size), self._api_max_batch)) all_embeddings: List[np.ndarray] = [] iterator = range(0, len(texts), step) if show_progress: try: from tqdm import tqdm iterator = tqdm(iterator, desc="Embedding (API)") except ImportError: pass for i in iterator: batch = texts[i : i + step] resp = self._client.embeddings.create( model=self.model, input=batch, encoding_format="float", ) rows = sorted( [(d.index, d.embedding) for d in resp.data], key=lambda x: x[0], ) batch_embs = np.array([e for _, e in rows], dtype=np.float32) if batch_embs.size > 0: self._set_dim_from_vector(batch_embs[0].tolist()) if normalize: norms = np.linalg.norm(batch_embs, axis=1, keepdims=True) batch_embs = batch_embs / (norms + 1e-10) all_embeddings.append(batch_embs) return np.vstack(all_embeddings).astype(np.float32) def similarity(self, emb1: np.ndarray, emb2: np.ndarray) -> np.ndarray: return np.dot(emb1, emb2.T) def encode_and_search( self, query: str, documents: List[str], top_k: int = 5, ) -> List[dict]: query_emb = self.encode([query], normalize=True) doc_embs = self.encode(documents, normalize=True) scores = self.similarity(query_emb, doc_embs)[0] top_indices = np.argsort(scores)[::-1][:top_k] return [ {"score": float(scores[idx]), "document": documents[idx], "index": int(idx)} for idx in top_indices ] # 向后兼容旧名称 DashScopeOpenAIEmbedding = OpenAICompatibleRemoteEmbedding _embedding_instance: Optional[Any] = None def get_embedder( _model_path: Optional[str] = None, _device: Optional[str] = None, force_reload: bool = False, ) -> OpenAICompatibleRemoteEmbedding: """ 获取远程 Embedding 单例。``_model_path`` / ``_device`` 已废弃，仅为兼容旧调用保留。 """ global _embedding_instance if force_reload or _embedding_instance is None: _embedding_instance = OpenAICompatibleRemoteEmbedding() return _embedding_instance # type: ignore[return-value] def clear_embedder() -> None: """清空单例（用于测试或切换远端配置）。""" global _embedding_instance _embedding_instance = None gc.collect()