sdpy-kit 0.1.0__py3-none-any.whl

kit/ai/profiles.py

@@ -0,0 +1,373 @@
+"""
+Copyright (c) 2021-2026 Clark Chang. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+Project: sdpy
+Author: Clark Chang
+"""
+
+# date: 2026-08-21
+"""模型档案（ModelProfile）与注册表（ModelRegistry）
+
+从 ``config/ai_models.json`` 读取四类型（chat / embedding / rerank / image）
+多厂商模型档案，用 pydantic 模型校验结构，并在 ``get()`` 时解析档案引用
+的环境变量密钥。
+
+密钥来源优先级：真实环境变量 > 项目根 ``.env``（``load()`` 时经
+``load_dotenv`` 注入，不覆盖已有环境变量）> 缺失报 ``AIConfigError``。
+档案数量与密钥变量名完全由 JSON 决定，无需改 ``sys_config``。
+"""
+
+import enum
+import os
+from dataclasses import dataclass, replace
+from functools import lru_cache
+from pathlib import Path
+from typing import Final
+
+from dotenv import load_dotenv
+from pydantic import BaseModel, Field, ValidationError
+
+from kit.ai.exceptions import AIConfigError
+from kit.config.kit_config import get_settings
+from kit.config.log_config import get_default_logger
+
+logger = get_default_logger(__name__)
+
+# 支持的模型类型
+_MODEL_TYPES: Final[tuple[str, ...]] = ("chat", "embedding", "rerank", "image")
+
+# 文件缺失报错时附带的最小可用模板（api_key 为环境变量名，密钥配置在 .env）
+_MINIMAL_MODELS_JSON: Final[str] = (
+    '{"chat": {"default_profile": "default", "profiles": {"default": '
+    '{"api_key": "SILICONFLOW_API_KEY", "base_url": "https://api.siliconflow.cn/v1", '
+    '"endpoint": "/chat/completions", "model": "Qwen/Qwen3-8B", "timeout": 120}}}}'
+)
+
+
+@dataclass(frozen=True)
+class ModelProfile:
+    """单个模型档案（运行时已解析密钥）。
+
+    Attributes:
+        name: 档案名（如 ``siliconflow``）
+        model_type: 模型类型（chat / embedding / rerank / image）
+        api_key_env: 密钥所在环境变量名
+        api_key: 已解析的密钥值
+        base_url: API 基础地址
+        endpoint: API 路径（以 ``/`` 开头）
+        model: 模型名称
+        timeout: 请求超时秒数
+        dimension: 向量维度（仅 embedding 有效）
+        batch_size: 单次请求最大文本数（仅 embedding 有效）
+    """
+
+    name: str
+    model_type: str
+    api_key_env: str
+    api_key: str
+    base_url: str
+    endpoint: str
+    model: str
+    timeout: int = 60
+    dimension: int | None = None
+    batch_size: int = 32
+
+    @property
+    def url(self) -> str:
+        """拼接完整请求 URL（base_url 去尾斜杠 + endpoint）"""
+        return self.base_url.rstrip("/") + self.endpoint
+
+
+class ProfileState(enum.Enum):
+    """档案运行时状态（候选链与状态快照使用）。"""
+
+    AVAILABLE = "available"  # 密钥已配置且可用（探测通过或尚未验证）
+    COOLING = "cooling"  # 调用失败冷却中（到期自动恢复）
+    INVALID_KEY = "invalid_key"  # 探测判定密钥无效
+    UNREACHABLE = "unreachable"  # 探测无法判定（超时 / 404 / 5xx）
+    UNCONFIGURED = "unconfigured"  # 密钥未配置
+    EXCLUDED = "excluded"  # 因约束被剔除（如 embedding 链内 model 不一致）
+
+
+@dataclass(frozen=True)
+class ProfileStatus:
+    """档案状态快照行（``ModelRegistry.available`` / ``AI.available`` 返回）。
+
+    Attributes:
+        name: 档案名
+        model: 模型名称
+        state: 运行时状态
+        current: 是否为该类型当前使用项（链头）
+        last_error: 最近一次失败原因（无失败为 None）
+        last_ok_at: 最近一次成功/探测通过时刻（Unix 秒；从未成功为 None）
+    """
+
+    name: str
+    model: str
+    state: ProfileState
+    current: bool = False
+    last_error: str | None = None
+    last_ok_at: float | None = None
+
+
+class ProfileConfig(BaseModel):
+    """档案配置（JSON 文件中单个 profile 的原始字段）。
+
+    注意：``api_key`` 存的是密钥所在环境变量名，
+    与 ``ModelProfile.api_key``（已解析的密钥值）同名不同义。
+    """
+
+    api_key: str = Field(min_length=1)
+    base_url: str = Field(min_length=1)
+    endpoint: str = Field(min_length=1)
+    model: str = Field(min_length=1)
+    timeout: int = Field(default=60, gt=0)
+    dimension: int | None = Field(default=None, gt=0)
+    batch_size: int = Field(default=32, gt=0)
+
+
+class TypeConfig(BaseModel):
+    """某模型类型的档案集合"""
+
+    default_profile: str = Field(min_length=1)
+    profiles: dict[str, ProfileConfig]
+
+
+class RegistryConfig(BaseModel):
+    """模型档案文件整体结构（顶层对象）"""
+
+    chat: TypeConfig | None = None
+    embedding: TypeConfig | None = None
+    rerank: TypeConfig | None = None
+    image: TypeConfig | None = None
+
+    def type_config(self, model_type: str) -> TypeConfig:
+        config = getattr(self, model_type, None)
+        if not isinstance(config, TypeConfig):
+            available = ", ".join(_MODEL_TYPES)
+            raise AIConfigError(f"模型类型不存在: {model_type}，可用类型: {available}")
+        return config
+
+
+class ModelRegistry:
+    """模型档案注册表。
+
+    用法::
+
+        registry = ModelRegistry()
+        registry.load("./config/ai_models.json")
+        profile = registry.get("chat", "siliconflow")
+    """
+
+    def __init__(self) -> None:
+        self._config: RegistryConfig | None = None
+
+    def load(self, path: str | None = None) -> None:
+        """加载并校验模型档案文件。
+
+        加载前将项目根 ``.env`` 注入 ``os.environ``（``override=False``：
+        真实环境变量优先，重复调用幂等），使密钥仅配置在 ``.env`` 时
+        ``get()`` 也能解析——档案数量与密钥变量名由 JSON 决定，
+        无需改 ``sys_config``。
+
+        Args:
+            path: 档案文件路径，None 时取 ``Settings.AI_MODELS_CONFIG_PATH``
+
+        Raises:
+            AIConfigError: 文件缺失、JSON 非法、结构非法、档案必填字段缺失、
+                默认档案未配置。
+        """
+        load_dotenv(Path.cwd() / ".env", override=False)
+
+        config_path = path or get_settings().AI_MODELS_CONFIG_PATH
+        file_path = Path(config_path)
+
+        if not file_path.is_file():
+            raise AIConfigError(
+                f"模型档案文件不存在: {config_path}。"
+                "请在项目根创建该文件（api_key 值为环境变量名，密钥配置在 .env），"
+                f"最小模板：\n{_MINIMAL_MODELS_JSON}\n"
+                "或调整 Settings.AI_MODELS_CONFIG_PATH 指向已有档案文件"
+            )
+
+        try:
+            self._config = RegistryConfig.model_validate_json(
+                file_path.read_text(encoding="utf-8")
+            )
+        except (OSError, ValidationError) as exc:
+            raise AIConfigError(f"模型档案文件解析失败: {config_path}: {exc}") from exc
+
+        loaded_types = [t for t in _MODEL_TYPES if getattr(self._config, t) is not None]
+        logger.info(
+            f"AI 模型档案已加载: {config_path}，类型: {', '.join(loaded_types) or '无'}"
+        )
+
+    def _resolve_profile(
+        self, model_type: str, profile: str | None
+    ) -> tuple[str, ProfileConfig]:
+        """解析类型 + 档案名，返回 (档案名, 静态 ProfileConfig)，不解析密钥。"""
+        if self._config is None:
+            raise AIConfigError("ModelRegistry 未加载，请先调用 load()")
+
+        type_config = self._config.type_config(model_type)
+        profile_name = profile or type_config.default_profile
+        profile_config = type_config.profiles.get(profile_name)
+        if profile_config is None:
+            raise AIConfigError(
+                f"档案不存在: {model_type}.{profile_name}，"
+                f"可用档案: {', '.join(sorted(type_config.profiles))}"
+            )
+        return profile_name, profile_config
+
+    def get(self, model_type: str, profile: str | None = None) -> ModelProfile:
+        """获取模型档案，profile 为 None 时回退到该类型默认档案。
+
+        Raises:
+            AIConfigError: 类型/档案不存在、密钥环境变量未配置。
+        """
+        profile_name, profile_config = self._resolve_profile(model_type, profile)
+
+        api_key = os.environ.get(profile_config.api_key) or ""
+        if not api_key:
+            raise AIConfigError(
+                f"未配置环境变量 {profile_config.api_key}，"
+                f"请检查 .env 中的 {profile_config.api_key}"
+            )
+
+        return self._build_profile(model_type, profile_name, profile_config, api_key)
+
+    @staticmethod
+    def _build_profile(
+        model_type: str, name: str, config: ProfileConfig, api_key: str
+    ) -> ModelProfile:
+        """由静态配置 + 已解析密钥构建运行时档案（get/chain 共用）。"""
+        return ModelProfile(
+            name=name,
+            model_type=model_type,
+            api_key_env=config.api_key,
+            api_key=api_key,
+            base_url=config.base_url,
+            endpoint=config.endpoint,
+            model=config.model,
+            timeout=config.timeout,
+            dimension=config.dimension,
+            batch_size=config.batch_size,
+        )
+
+    def get_profile(self, model_type: str, profile: str | None = None) -> ProfileConfig:
+        """读取静态档案（不含密钥），用于维度/模型名等无需密钥的配置。"""
+        return self._resolve_profile(model_type, profile)[1]
+
+    def type_profiles(self, model_type: str) -> dict[str, ProfileConfig]:
+        """返回该类型全部静态档案（含密钥未配置者），供状态快照展示。
+
+        Raises:
+            AIConfigError: 注册表未加载或类型不存在。
+        """
+        if self._config is None:
+            raise AIConfigError("ModelRegistry 未加载，请先调用 load()")
+        return dict(self._config.type_config(model_type).profiles)
+
+    def chain(
+        self, model_type: str, *, env_profiles: str | None = None
+    ) -> list[ModelProfile]:
+        """解析该类型候选链：优先级链（env）> default 打头 + JSON 声明顺序。
+
+        规则：
+            1. ``env_profiles`` 参数或 ``Settings.AI_{TYPE}_PROFILES``（逗号
+               分隔档案名）设置时按其顺序；JSON 中不存在的档案名丢弃并告警；
+            2. 未设置时 ``default_profile`` 打头，其余按 JSON 声明顺序；
+            3. 密钥未配置的档案剔除出链（debug 日志）。
+
+        Args:
+            model_type: 模型类型（chat / embedding / rerank / image）
+            env_profiles: 显式优先级链（测试注入用）；None 时读 Settings
+
+        Returns:
+            候选档案列表（密钥已解析；可能为空——全部未配置密钥）。
+
+        Raises:
+            AIConfigError: 注册表未加载或类型不存在。
+        """
+        if self._config is None:
+            raise AIConfigError("ModelRegistry 未加载，请先调用 load()")
+        type_config = self._config.type_config(model_type)
+
+        env_value = env_profiles
+        if env_value is None:
+            settings = get_settings()
+            env_value = {
+                "chat": settings.AI_CHAT_PROFILES,
+                "embedding": settings.AI_EMBEDDING_PROFILES,
+                "rerank": settings.AI_RERANK_PROFILES,
+                "image": settings.AI_IMAGE_PROFILES,
+            }.get(model_type)
+
+        if env_value:
+            names = [name.strip() for name in env_value.split(",") if name.strip()]
+            unknown = [name for name in names if name not in type_config.profiles]
+            if unknown:
+                logger.warning(
+                    f"{model_type} 候选链忽略 JSON 中不存在的档案: "
+                    f"{', '.join(unknown)}"
+                )
+            order = [name for name in names if name in type_config.profiles]
+        else:
+            order = [
+                type_config.default_profile,
+                *(
+                    name
+                    for name in type_config.profiles
+                    if name != type_config.default_profile
+                ),
+            ]
+
+        profiles: list[ModelProfile] = []
+        for name in order:
+            config = type_config.profiles[name]
+            api_key = os.environ.get(config.api_key) or ""
+            if not api_key:
+                logger.debug(
+                    f"{model_type}.{name} 密钥未配置（{config.api_key}），剔除出链"
+                )
+                continue
+            profiles.append(self._build_profile(model_type, name, config, api_key))
+        return profiles
+
+
+def with_api_key(profile: ModelProfile, api_key: str) -> ModelProfile:
+    """返回替换 api_key 后的新档案（ModelProfile 为不可变）。"""
+    return replace(profile, api_key=api_key)
+
+
+@lru_cache(maxsize=1)
+def get_model_registry() -> ModelRegistry:
+    """获取全局单例 ModelRegistry（首次调用自动 load 默认档案文件）。"""
+    registry = ModelRegistry()
+    registry.load()
+    return registry
+
+
+def embedding_dimension() -> int:
+    """embedding 默认档案的向量维度（读静态字段，无需密钥）。"""
+    dimension = get_model_registry().get_profile("embedding").dimension
+    if dimension is None:
+        raise AIConfigError("ai_models.json 的 embedding 档案缺少 dimension")
+    return dimension
+
+
+def embedding_model() -> str:
+    """embedding 默认档案的模型名（读静态字段，无需密钥）。"""
+    return get_model_registry().get_profile("embedding").model

kit/ai/rerank_service.py

@@ -0,0 +1,93 @@
+"""
+Copyright (c) 2021-2026 Clark Chang. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+Project: sdpy
+Author: Clark Chang
+"""
+
+from dataclasses import dataclass
+from typing import cast, override
+
+from kit.ai.base_service import BaseAIService, safe_get
+from kit.ai.exceptions import AIResponseError
+from kit.ai.types import JSONObject
+
+__all__ = ["RerankRequest", "RerankHit", "RerankService"]
+
+
+@dataclass(frozen=True)
+class RerankRequest:
+    """重排序请求模型。"""
+
+    query: str
+    documents: list[str]
+    top_n: int | None = None
+
+
+@dataclass(frozen=True)
+class RerankHit:
+    """重排序单条命中结果。"""
+
+    index: int
+    score: float
+
+
+class RerankService(BaseAIService[RerankRequest, list[RerankHit]]):
+    """重排序服务（全异步）。"""
+
+    @override
+    def _build_payload(self, request: RerankRequest) -> JSONObject:
+        payload: dict[str, object] = {
+            "model": self._profile.model,
+            "query": request.query,
+            "documents": request.documents,
+        }
+        if request.top_n is not None:
+            payload["top_n"] = request.top_n
+        return cast(JSONObject, payload)
+
+    @override
+    def _parse_response(self, data: JSONObject) -> list[RerankHit]:
+        items = safe_get(data, "results")
+        if not isinstance(items, list):
+            raise AIResponseError(f"重排序响应 results 必须是数组: {str(data)[:200]}")
+        hits: list[RerankHit] = []
+        for item in items:
+            if not isinstance(item, dict):
+                raise AIResponseError(f"重排序响应条目非法: {str(item)[:200]}")
+            index = item.get("index")
+            score = item.get("relevance_score")
+            if not isinstance(index, int) or not isinstance(score, (int, float)):
+                raise AIResponseError(f"重排序响应字段异常: {str(item)[:200]}")
+            hits.append(RerankHit(index=index, score=float(score)))
+        hits.sort(key=lambda h: h.score, reverse=True)
+        return hits
+
+    async def rerank(
+        self,
+        query: str,
+        documents: list[str],
+        *,
+        top_n: int | None = None,
+    ) -> list[RerankHit]:
+        if not documents:
+            return []
+        hits = await self._call(RerankRequest(query, documents, top_n=top_n))
+        for hit in hits:
+            if hit.index < 0 or hit.index >= len(documents):
+                raise AIResponseError(
+                    f"重排序 index 越界: {hit.index} (文档数 {len(documents)})"
+                )
+        return hits