"""DashScope / Qwen 远程嵌入实现。 通过 OpenAI 兼容接口调用阿里百炼的 text-embedding-v3: base_url: https://dashscope.aliyuncs.com/compatible-mode/v1 model: text-embedding-v3 (1024 维) 限制: 单次请求 input ≤ 25 条 配置来源(优先级从高到低): 1. 构造参数(model / api_key / base_url / max_attempts) 2. configs/llm_models.yaml 的 scenes.embedding 3. 环境变量 / .env: DASHSCOPE_EMBEDDING_API_KEY / DASHSCOPE_API_KEY DASHSCOPE_EMBEDDING_BASE_URL / QWEN_BASE_URL DASHSCOPE_EMBEDDING_MODEL 4. 代码内置默认值(text-embedding-v3 / 1024 维) """ from __future__ import annotations import asyncio from loguru import logger from openai import AsyncOpenAI, OpenAI from configs.loader import load_scene_config from configs.runtime_env import env_get from .base import AsyncEmbeddingProvider, EmbeddingProvider from .models import EmbeddingError DASHSCOPE_DEFAULT_BASE = "https://dashscope.aliyuncs.com/compatible-mode/v1" DASHSCOPE_DEFAULT_MODEL = "text-embedding-v3" DASHSCOPE_DEFAULT_DIM = 1024 DASHSCOPE_BATCH_LIMIT = 10 # 百炼实测单批上限(2026-06,文档曾标 25 但 API 报 400) # 重试策略 DEFAULT_MAX_ATTEMPTS = 3 RETRY_BASE_WAIT_SEC = 1.0 RETRY_MAX_WAIT_SEC = 8.0 # embedding 场景名(对应 configs/llm_models.yaml scenes.embedding) SCENE_EMBEDDING = "embedding" def _read_env(key: str, default: str | None = None) -> str | None: """读取环境变量（先热加载 .env，改文件后无需重启进程）。""" return env_get(key, default) def _scene() -> dict: """读取 YAML embedding 场景配置(不存在时为空 dict)。""" return load_scene_config(SCENE_EMBEDDING) def _scene_int(key: str, default: int) -> int: try: return int(_scene().get(key) or default) except (TypeError, ValueError): return default def _scene_float(key: str, default: float) -> float: try: return float(_scene().get(key) or default) except (TypeError, ValueError): return default def _resolve_config() -> tuple[str, str, str]: """读取 API key / base_url / model,返回 (api_key, base_url, model)。 优先级: YAML 场景 > 环境变量 > 内置默认。 模型名优先级:DASHSCOPE_EMBEDDING_MODEL > 默认值。 不再读全局 EMBEDDING_MODEL,避免与 LOCAL provider 冲突。 """ sc = _scene() api_key_env = sc.get("api_key_env") or "DASHSCOPE_EMBEDDING_API_KEY" api_key = _read_env(api_key_env) or _read_env("DASHSCOPE_API_KEY") or "" if not api_key: raise EmbeddingError( f"{api_key_env} 或 DASHSCOPE_API_KEY 未配置" ) # YAML base_url_env -> DASHSCOPE_EMBEDDING_BASE_URL -> QWEN_BASE_URL(兜底) -> 默认 base_url = ( _read_env(sc.get("base_url_env") or "DASHSCOPE_EMBEDDING_BASE_URL") or _read_env("QWEN_BASE_URL") or DASHSCOPE_DEFAULT_BASE ) model = ( sc.get("model") or _read_env("DASHSCOPE_EMBEDDING_MODEL", DASHSCOPE_DEFAULT_MODEL) or DASHSCOPE_DEFAULT_MODEL ) return api_key, base_url, model def _chunked(items: list[str], size: int) -> list[list[str]]: """把列表按 size 分块。""" return [items[i : i + size] for i in range(0, len(items), size)] class DashScopeEmbeddingProvider(EmbeddingProvider): """同步实现,主要用于测试/单条调用。""" name = "dashscope" def __init__( self, *, model: str | None = None, api_key: str | None = None, base_url: str | None = None, timeout_sec: float | None = None, max_attempts: int | None = None, batch_limit: int | None = None, ) -> None: env_key, env_base, env_model = _resolve_config() self.model = model or env_model self.dim = DASHSCOPE_DEFAULT_DIM self.max_attempts = max_attempts or _scene_int("max_attempts", DEFAULT_MAX_ATTEMPTS) self.batch_limit = batch_limit or _scene_int("batch_limit", DASHSCOPE_BATCH_LIMIT) timeout = timeout_sec or _scene_float("timeout_sec", 60.0) self._client = OpenAI( api_key=api_key or env_key, base_url=base_url or env_base, timeout=timeout, ) def embed_batch(self, texts: list[str]) -> list[list[float]]: if not texts: return [] results: list[list[float]] = [] for chunk in _chunked(texts, self.batch_limit): results.extend(self._call_with_retry(chunk)) return results def _call_with_retry(self, batch: list[str]) -> list[list[float]]: import time last_err: Exception | None = None for attempt in range(1, self.max_attempts + 1): try: resp = self._client.embeddings.create(model=self.model, input=batch) return [d.embedding for d in resp.data] except Exception as e: # noqa: BLE001 last_err = e logger.warning( "DashScope embed 失败 尝试 {}/{}: {}: {}", attempt, self.max_attempts, type(e).__name__, e, ) if attempt < self.max_attempts: wait = min(RETRY_BASE_WAIT_SEC * (2 ** (attempt - 1)), RETRY_MAX_WAIT_SEC) time.sleep(wait) raise EmbeddingError( f"DashScope embed 放弃 {self.max_attempts} 次: {last_err}", attempts=self.max_attempts, ) def close(self) -> None: self._client.close() class DashScopeAsyncEmbeddingProvider(AsyncEmbeddingProvider): """异步实现,用于批处理。""" name = "dashscope" def __init__( self, *, model: str | None = None, api_key: str | None = None, base_url: str | None = None, timeout_sec: float | None = None, max_attempts: int | None = None, batch_limit: int | None = None, ) -> None: env_key, env_base, env_model = _resolve_config() self.model = model or env_model self.dim = DASHSCOPE_DEFAULT_DIM self.max_attempts = max_attempts or _scene_int("max_attempts", DEFAULT_MAX_ATTEMPTS) self.batch_limit = batch_limit or _scene_int("batch_limit", DASHSCOPE_BATCH_LIMIT) timeout = timeout_sec or _scene_float("timeout_sec", 60.0) self._client = AsyncOpenAI( api_key=api_key or env_key, base_url=base_url or env_base, timeout=timeout, ) async def embed_batch(self, texts: list[str]) -> list[list[float]]: if not texts: return [] results: list[list[float]] = [] for chunk in _chunked(texts, self.batch_limit): results.extend(await self._call_with_retry(chunk)) return results async def _call_with_retry(self, batch: list[str]) -> list[list[float]]: last_err: Exception | None = None for attempt in range(1, self.max_attempts + 1): try: resp = await self._client.embeddings.create( model=self.model, input=batch ) return [d.embedding for d in resp.data] except Exception as e: # noqa: BLE001 last_err = e logger.warning( "DashScope embed 失败 尝试 {}/{}: {}: {}", attempt, self.max_attempts, type(e).__name__, e, ) if attempt < self.max_attempts: wait = min(RETRY_BASE_WAIT_SEC * (2 ** (attempt - 1)), RETRY_MAX_WAIT_SEC) await asyncio.sleep(wait) raise EmbeddingError( f"DashScope embed 放弃 {self.max_attempts} 次: {last_err}", attempts=self.max_attempts, ) async def close(self) -> None: await self._client.close()