初始化

2026-07-18 16:13:52 +08:00
parent c0070f0a5c
commit fe8b417ab6
75 changed files with 12898 additions and 1 deletions
@@ -0,0 +1,198 @@
+"""DashScope Embedding 客户端。
+
+通过 OpenAI 兼容接口调用阿里百炼 text-embedding-v3：
+    base_url: https://dashscope.aliyuncs.com/compatible-mode/v1
+    model:    text-embedding-v3（1024 维）
+
+配置来源:
+    - .env → DASHSCOPE_API_KEY / QWEN_BASE_URL（OpenAI 兼容端点）
+    - configs/system.yaml → embedding 段（model / dimension / batch_size / timeout）
+"""
+
+import logging
+import os
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+import yaml
+from openai import OpenAI
+
+from embedding.models import EmbeddingError
+
+logger = logging.getLogger(__name__)
+
+# 默认值
+DASHSCOPE_DEFAULT_BASE = "https://dashscope.aliyuncs.com/compatible-mode/v1"
+DASHSCOPE_DEFAULT_MODEL = "text-embedding-v3"
+DASHSCOPE_DEFAULT_DIM = 1024
+DASHSCOPE_BATCH_LIMIT = 10  # 百炼实测单批上限
+
+# 重试
+DEFAULT_MAX_ATTEMPTS = 3
+RETRY_BASE_WAIT_SEC = 1.0
+RETRY_MAX_WAIT_SEC = 8.0
+
+
+def _load_embedding_config() -> dict:
+    """从 system.yaml 加载 embedding 段配置。"""
+    config_path = Path("configs/system.yaml")
+    if config_path.exists():
+        try:
+            with open(config_path, encoding="utf-8") as f:
+                raw = yaml.safe_load(f)
+            return raw.get("embedding", {})
+        except Exception:
+            logger.warning("加载 embedding 配置失败")
+    return {}
+
+
+@dataclass
+class EmbeddingConfig:
+    """Embedding 调用配置。"""
+
+    provider: str = "dashscope"
+    model: str = DASHSCOPE_DEFAULT_MODEL
+    api_key: str = ""
+    base_url: str = DASHSCOPE_DEFAULT_BASE
+    dimension: int = DASHSCOPE_DEFAULT_DIM
+    batch_size: int = DASHSCOPE_BATCH_LIMIT
+    timeout_sec: float = 30.0
+    max_attempts: int = DEFAULT_MAX_ATTEMPTS
+
+    def __post_init__(self) -> None:
+        if not self.api_key:
+            raise EmbeddingError("DASHSCOPE_API_KEY 未配置，请检查 .env")
+
+
+def load_embedding_config(
+    *,
+    model: str | None = None,
+) -> EmbeddingConfig:
+    """根据配置构造 EmbeddingConfig。
+
+    优先级: system.yaml > .env 默认值 > 硬编码默认值
+    """
+    sys_cfg = _load_embedding_config()
+
+    # API key: 从环境变量读取
+    api_key = os.environ.get("DASHSCOPE_API_KEY", "")
+    if not api_key:
+        api_key = os.environ.get("QWEN_API_KEY", "")
+
+    # Base URL: 优先用 QWEN_BASE_URL（OpenAI 兼容），
+    # DASHSCOPE_BASE_URL 通常是旧版非兼容端点，不作为默认
+    base_url = (
+        os.environ.get("QWEN_BASE_URL")
+        or os.environ.get("DASHSCOPE_EMBEDDING_BASE_URL")
+        or DASHSCOPE_DEFAULT_BASE
+    )
+
+    m = model or sys_cfg.get("dashscope_model", DASHSCOPE_DEFAULT_MODEL)
+    dimension = int(sys_cfg.get("dimension", DASHSCOPE_DEFAULT_DIM))
+    batch_size = min(int(sys_cfg.get("batch_size", DASHSCOPE_BATCH_LIMIT)), DASHSCOPE_BATCH_LIMIT)
+    timeout = float(sys_cfg.get("timeout_sec", 30.0))
+
+    if not api_key:
+        raise EmbeddingError("DASHSCOPE_API_KEY 未配置，请检查 .env")
+
+    return EmbeddingConfig(
+        provider="dashscope",
+        model=m,
+        api_key=api_key,
+        base_url=base_url,
+        dimension=dimension,
+        batch_size=batch_size,
+        timeout_sec=timeout,
+    )
+
+
+def make_embedding_client(config: EmbeddingConfig) -> OpenAI:
+    """构造同步 OpenAI 客户端（指向 DashScope 兼容端点）。"""
+    logger.info(
+        "初始化 Embedding 客户端: provider=%s model=%s base_url=%s dim=%d",
+        config.provider, config.model, config.base_url, config.dimension,
+    )
+    return OpenAI(
+        api_key=config.api_key,
+        base_url=config.base_url,
+        timeout=config.timeout_sec,
+    )
+
+
+def _chunked(items: list[str], size: int) -> list[list[str]]:
+    """把列表按 size 分块。"""
+    return [items[i : i + size] for i in range(0, len(items), size)]
+
+
+def embed_batch(
+    client: OpenAI,
+    config: EmbeddingConfig,
+    texts: list[str],
+) -> list[list[float]]:
+    """批量嵌入，自动分块+重试。
+
+    Args:
+        client: OpenAI 客户端
+        config: Embedding 配置
+        texts: 待嵌入文本列表
+
+    Returns:
+        与 texts 等长的向量列表，每个为 1024 维 float 列表
+    """
+    if not texts:
+        return []
+
+    all_results: list[list[float]] = []
+    chunks = _chunked(texts, config.batch_size)
+
+    for chunk_idx, chunk in enumerate(chunks):
+        result = _call_with_retry(client, config, chunk, chunk_idx, len(chunks))
+        all_results.extend(result)
+
+    return all_results
+
+
+def _call_with_retry(
+    client: OpenAI,
+    config: EmbeddingConfig,
+    batch: list[str],
+    chunk_idx: int,
+    total_chunks: int,
+) -> list[list[float]]:
+    """单批嵌入调用，带指数退避重试。"""
+    last_err: Exception | None = None
+
+    for attempt in range(1, config.max_attempts + 1):
+        try:
+            resp = client.embeddings.create(model=config.model, input=batch)
+            vectors = [d.embedding for d in resp.data]
+
+            # 维度校验
+            if vectors and len(vectors[0]) != config.dimension:
+                logger.warning(
+                    "实际维度 %d 与预期 %d 不一致",
+                    len(vectors[0]), config.dimension,
+                )
+
+            logger.debug(
+                "Embedding chunk %d/%d 完成（%d 条，attempt %d）",
+                chunk_idx + 1, total_chunks, len(batch), attempt,
+            )
+            return vectors
+
+        except Exception as e:
+            last_err = e
+            logger.warning(
+                "DashScope embed 失败 chunk %d/%d 尝试 %d/%d: %s: %s",
+                chunk_idx + 1, total_chunks, attempt, config.max_attempts,
+                type(e).__name__, e,
+            )
+            if attempt < config.max_attempts:
+                wait = min(RETRY_BASE_WAIT_SEC * (2 ** (attempt - 1)), RETRY_MAX_WAIT_SEC)
+                time.sleep(wait)
+
+    raise EmbeddingError(
+        f"DashScope embed 放弃（{config.max_attempts} 次）: {last_err}",
+        attempts=config.max_attempts,
+    )