news/embedding/remote.py

"""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()