sdpy-kit 0.1.0__py3-none-any.whl

kit/ai/pool.py

@@ -0,0 +1,411 @@
+"""
+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: SynthArk
+Author: Clark Chang
+"""
+
+# date: 2026-09-05
+"""AI 服务候选池：激活探测、故障冷却、当前项切换
+
+按模型类型维护「候选链」：链上档案按优先级排序（``prefer()`` 指定项打头），
+仅 ``AVAILABLE`` 状态档案参与调用竞选；档案调用失败进入冷却（默认 60 秒），
+冷却期内被跳过（备胎顶上），到期自动恢复——冷却即粘性降级，自愈无需干预。
+
+预热探测（``probe_all``）经 ``GET {base_url}/models`` 验证密钥有效性：
+2xx / 429 视为 key 有效并激活；401 / 403 判定 key 无效出链；其余
+（超时 / 404 / 5xx）无法判定出链。剔除后链空时回退「密钥已配置档案全部
+惰性入链」并告警，保证功能不瘫。
+
+线程模型：业务调用与探测均在 runner 后台事件循环线程内执行；
+``available()`` / ``prefer()`` 可能从业务线程调用，内部以锁保护共享状态。
+"""
+
+import asyncio
+import threading
+import time
+from collections.abc import Callable, Iterator
+from dataclasses import dataclass, field
+from typing import Final
+
+from kit.ai.chat_service import ChatService
+from kit.ai.embedding_service import EmbeddingService
+from kit.ai.exceptions import AIConfigError, AIError, AIRequestError
+from kit.ai.image_service import ImageService
+from kit.ai.model_builder import Service
+from kit.ai.model_client import ModelClient
+from kit.ai.profiles import (
+    ModelProfile,
+    ModelRegistry,
+    ProfileState,
+    ProfileStatus,
+    get_model_registry,
+)
+from kit.ai.rerank_service import RerankService
+from kit.config.log_config import get_default_logger
+
+__all__ = ["AIServicePool"]
+
+logger = get_default_logger(__name__)
+
+# 调用失败后的冷却时长（秒）：冷却期内备胎顶上，到期自动回位
+_COOLDOWN_SECONDS: Final = 60.0
+
+# 各模型类型 → 单档案服务工厂
+_FACTORIES: Final[dict[str, Callable[[ModelProfile], Service]]] = {
+    "chat": ChatService,
+    "embedding": EmbeddingService,
+    "rerank": RerankService,
+    "image": ImageService,
+}
+
+# 探测判定为「密钥无效」的 HTTP 状态码（2xx 已由 list_models 放行；
+# 429 限流恰证明服务与密钥均有效，在 _probe_state 中特判为可用）
+_PROBE_INVALID_STATUS: Final[frozenset[int]] = frozenset({401, 403})
+
+
+async def _probe_profile(profile: ModelProfile) -> None:
+    """单档案探测：key 有效（2xx / 429）静默返回，否则抛 ``AIError``。
+
+    独立模块级函数便于测试替换。探测客户端用毕即关。
+    """
+    client = ModelClient(profile)
+    try:
+        await client.list_models()
+    finally:
+        await client.close()
+
+
+def _probe_state(exc: BaseException) -> ProfileState:
+    """探测异常 → 档案状态：429 视为 key 有效；401/403 无效；其余无法判定。"""
+    if isinstance(exc, AIRequestError) and exc.status_code == 429:
+        return ProfileState.AVAILABLE
+    if (
+        isinstance(exc, AIRequestError)
+        and exc.status_code in _PROBE_INVALID_STATUS
+    ):
+        return ProfileState.INVALID_KEY
+    return ProfileState.UNREACHABLE
+
+
+class _Entry:
+    """链内单档案的运行时状态（可变）。"""
+
+    __slots__ = (
+        "profile",
+        "state",
+        "cooling_until",
+        "last_error",
+        "last_ok_at",
+        "service",
+    )
+
+    def __init__(self, profile: ModelProfile) -> None:
+        self.profile: ModelProfile = profile
+        self.state: ProfileState = ProfileState.AVAILABLE
+        # 冷却截止时刻（time.monotonic 秒）；0 表示未冷却
+        self.cooling_until: float = 0.0
+        self.last_error: str | None = None
+        self.last_ok_at: float | None = None
+        self.service: Service | None = None
+
+
+@dataclass
+class _Chain:
+    """某类型的候选链：链序 + 全档案状态表（entries 含被剔除档案）。"""
+
+    order: list[str] = field(default_factory=list)
+    entries: dict[str, _Entry] = field(default_factory=dict)
+
+
+class AIServicePool:
+    """按模型类型维护候选链：状态表 + 冷却 + 当前项 + service 惰性单例。"""
+
+    def __init__(self, registry: ModelRegistry | None = None) -> None:
+        self._registry: ModelRegistry = registry or get_model_registry()
+        self._lock = threading.Lock()
+        self._chains: dict[str, _Chain] = {}
+        # prefer() 指定的长期当前项（类型 → 档案名）；进程重启即失效
+        self._preferred: dict[str, str] = {}
+
+    # ---------- 链解析 ----------
+
+    def _ensure_chain(self, model_type: str) -> _Chain:
+        """取（或首次解析）类型候选链。
+
+        ``registry.chain`` 涉及读配置与环境变量，置于锁外执行；
+        锁内双检避免并发重复解析。
+        """
+        with self._lock:
+            chain = self._chains.get(model_type)
+        if chain is not None:
+            return chain
+        profiles = self._registry.chain(model_type)
+        with self._lock:
+            chain = self._chains.get(model_type)
+            if chain is not None:
+                return chain
+            entries: dict[str, _Entry] = {}
+            order: list[str] = []
+            for profile in profiles:
+                if profile.name not in entries:
+                    entries[profile.name] = _Entry(profile)
+                order.append(profile.name)
+            # embedding 链内强制同 model：维度一致性是向量库可用性的前提
+            if model_type == "embedding" and order:
+                base_model = entries[order[0]].profile.model
+                excluded = [
+                    name
+                    for name in order
+                    if entries[name].profile.model != base_model
+                ]
+                for name in excluded:
+                    entries[name].state = ProfileState.EXCLUDED
+                    entries[name].last_error = (
+                        f"embedding 链内 model 不一致（要求 {base_model}）"
+                    )
+                    order.remove(name)
+                if excluded:
+                    logger.warning(
+                        "embedding 候选链剔除 model 不一致档案: "
+                        f"{', '.join(excluded)}（链内仅允许 {base_model}）"
+                    )
+            chain = _Chain(order=order, entries=entries)
+            self._chains[model_type] = chain
+            return chain
+
+    # ---------- 探测 ----------
+
+    async def probe_all(self) -> None:
+        """后台预热探测：逐类型并行验证候选档案密钥有效性（不抛异常）。
+
+        单类型未配置（json 缺该类型）仅跳过，不影响其余类型探测。
+        """
+        for model_type in _FACTORIES:
+            try:
+                await self._probe_type(model_type)
+            except AIConfigError as exc:
+                logger.debug(f"{model_type} 探测跳过（类型未配置档案）: {exc}")
+            except Exception as exc:  # noqa: BLE001 - 预热失败不影响主流程
+                logger.warning(f"{model_type} 档案探测异常: {exc}")
+
+    async def _probe_type(self, model_type: str) -> None:
+        """探测单类型链上档案，按判定表激活或剔除；链空时回退惰性入链。"""
+        chain = self._ensure_chain(model_type)
+        with self._lock:
+            names = list(chain.order)
+        if not names:
+            return
+        results = await asyncio.gather(
+            *(_probe_profile(chain.entries[name].profile) for name in names),
+            return_exceptions=True,
+        )
+        dropped: list[str] = []
+        with self._lock:
+            for name, result in zip(names, results, strict=True):
+                entry = chain.entries[name]
+                if isinstance(result, BaseException):
+                    state = _probe_state(result)
+                    entry.state = state
+                    entry.last_error = f"探测: {result}"
+                    if state is ProfileState.AVAILABLE:
+                        # 429：限流但 key 有效，保持可用
+                        entry.last_ok_at = time.time()
+                    else:
+                        dropped.append(name)
+                        logger.warning(
+                            f"{model_type}.{name} 探测未通过，剔除出链: {result}"
+                        )
+                else:
+                    entry.state = ProfileState.AVAILABLE
+                    entry.last_ok_at = time.time()
+                    entry.last_error = None
+            if dropped:
+                for name in dropped:
+                    if name in chain.order:
+                        chain.order.remove(name)
+                # 兜底：剔除后链空 → 回退「密钥已配置档案全部惰性入链」
+                if not chain.order:
+                    chain.order = list(names)
+                    for name in names:
+                        chain.entries[name].state = ProfileState.AVAILABLE
+                    logger.warning(
+                        f"{model_type} 探测后候选链为空，回退惰性入链: "
+                        f"{', '.join(names)}（真实调用时再验证）"
+                    )
+                else:
+                    logger.info(
+                        f"{model_type} 探测完成: 可用 "
+                        f"{', '.join(chain.order)}；剔除 {', '.join(dropped)}"
+                    )
+
+    # ---------- 候选迭代 ----------
+
+    def iter_services(self, model_type: str) -> Iterator[tuple[str, Service]]:
+        """按竞选顺序迭代（档案名, service）：prefer 项打头，跳过冷却档案。
+
+        候选计划（顺序 + prefer 提前 + 冷却到期复位 + 可用性过滤）在生成器
+        首次推进时一次定格：此后 ``prefer()`` 切换与其它请求的失败冷却均
+        不影响已开始迭代的调用方——切换只为「下一次」请求生效。
+
+        仅在后台事件循环线程内迭代；service 惰性构建并缓存复用。
+        """
+        chain = self._ensure_chain(model_type)
+        now = time.monotonic()
+        with self._lock:
+            seq = list(chain.order)
+            preferred = self._preferred.get(model_type)
+            if preferred is not None and preferred in seq:
+                seq.remove(preferred)
+                seq.insert(0, preferred)
+            plan: list[tuple[str, Service | None]] = []
+            for name in seq:
+                entry = chain.entries[name]
+                if (
+                    entry.state is ProfileState.COOLING
+                    and now >= entry.cooling_until
+                ):
+                    entry.state = ProfileState.AVAILABLE
+                    entry.last_error = None
+                if entry.state is not ProfileState.AVAILABLE:
+                    continue
+                plan.append((name, entry.service))
+        for name, service in plan:
+            if service is None:
+                entry = chain.entries[name]
+                service = _FACTORIES[model_type](entry.profile)
+                with self._lock:
+                    if entry.service is None:
+                        entry.service = service
+                    else:
+                        service = entry.service
+            yield name, service
+
+    # ---------- 状态标记 ----------
+
+    def mark_failure(self, model_type: str, name: str, exc: AIError) -> None:
+        """档案调用失败：进入冷却并记录原因。"""
+        with self._lock:
+            chain = self._chains.get(model_type)
+            entry = chain.entries.get(name) if chain else None
+            if entry is None:
+                return
+            entry.state = ProfileState.COOLING
+            entry.cooling_until = time.monotonic() + _COOLDOWN_SECONDS
+            entry.last_error = f"{type(exc).__name__}: {exc}"
+        logger.warning(
+            f"{model_type}.{name} 调用失败，冷却 {_COOLDOWN_SECONDS:g}s: {exc}"
+        )
+
+    def mark_success(self, model_type: str, name: str) -> None:
+        """档案调用成功：记录成功时刻并清除失败痕迹。"""
+        with self._lock:
+            chain = self._chains.get(model_type)
+            entry = chain.entries.get(name) if chain else None
+            if entry is None:
+                return
+            entry.last_ok_at = time.time()
+            entry.last_error = None
+
+    # ---------- 当前项与状态 ----------
+
+    def prefer(self, model_type: str, profile: str) -> None:
+        """切换该类型长期当前项：指定档案成为固定链头并清除冷却。
+
+        Raises:
+            AIConfigError: 档案不存在，或不可用
+                （未配置 / 密钥无效 / 不可达 / 因约束被排除）。
+        """
+        chain = self._ensure_chain(model_type)
+        with self._lock:
+            entry = chain.entries.get(profile)
+            if entry is None:
+                known = ", ".join(chain.entries) or "无"
+                raise AIConfigError(
+                    f"档案不存在: {model_type}.{profile}（已知档案: {known}）"
+                )
+            if entry.state in (
+                ProfileState.UNCONFIGURED,
+                ProfileState.INVALID_KEY,
+                ProfileState.UNREACHABLE,
+                ProfileState.EXCLUDED,
+            ):
+                raise AIConfigError(
+                    f"档案不可用: {model_type}.{profile}"
+                    f"（状态 {entry.state.value}，原因: {entry.last_error}）"
+                )
+            entry.state = ProfileState.AVAILABLE
+            entry.cooling_until = 0.0
+            self._preferred[model_type] = profile
+        logger.info(f"已切换 {model_type} 当前使用档案: {profile}")
+
+    def available(self, model_type: str) -> list[ProfileStatus]:
+        """该类型全部档案状态快照（含未配置与被剔除者），链头标记 current。"""
+        chain = self._ensure_chain(model_type)
+        now = time.monotonic()
+        with self._lock:
+            preferred = self._preferred.get(model_type)
+            if preferred in chain.order:
+                current: str | None = preferred
+            else:
+                current = chain.order[0] if chain.order else None
+            rows = [
+                ProfileStatus(
+                    name=name,
+                    model=entry.profile.model,
+                    state=(
+                        ProfileState.AVAILABLE
+                        if entry.state is ProfileState.COOLING
+                        and now >= entry.cooling_until
+                        else entry.state
+                    ),
+                    current=name == current,
+                    last_error=entry.last_error,
+                    last_ok_at=entry.last_ok_at,
+                )
+                for name, entry in chain.entries.items()
+            ]
+        try:
+            all_profiles = self._registry.type_profiles(model_type)
+        except AIConfigError:
+            all_profiles = {}
+        for name, config in all_profiles.items():
+            if name not in chain.entries:
+                rows.append(
+                    ProfileStatus(
+                        name=name,
+                        model=config.model,
+                        state=ProfileState.UNCONFIGURED,
+                        current=False,
+                    )
+                )
+        return rows
+
+    # ---------- 生命周期 ----------
+
+    async def close_all(self) -> None:
+        """关闭全部已构建 service（释放底层 httpx 连接）。幂等。"""
+        with self._lock:
+            services = [
+                entry.service
+                for chain in self._chains.values()
+                for entry in chain.entries.values()
+                if entry.service is not None
+            ]
+        for service in services:
+            try:
+                await service.close()
+            except Exception as exc:  # noqa: BLE001 - 关闭阶段失败不影响退出
+                logger.warning(f"关闭 {service} 失败: {exc}")