sdpy-kit 0.1.0__py3-none-any.whl

kit/ai/failover.py

@@ -0,0 +1,244 @@
+"""
+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
+"""候选链故障切换调用层
+
+将单档案 Service 的公开方法包装为链式调用：按池内候选顺序尝试，任一档案
+抛出 ``AIError`` 即冷却该档案并尝试下一个；全部失败抛出携带已尝试档案
+列表与最后失败原因的 ``AIRequestError``。
+
+流式对话仅在「首帧到达前」允许切换档案；首帧已产出后失败直接抛出
+（已输出 token 无法回滚）。空流（首帧即结束）视为正常完成。
+
+各 Failover 服务公开方法与对应单档案 Service 保持同形，业务侧无感替换。
+"""
+
+from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable
+from typing import TypeVar, cast
+
+from kit.ai.chat_service import ChatService
+from kit.ai.embedding_service import EmbeddingService
+from kit.ai.exceptions import AIError, AIRequestError
+from kit.ai.image_service import ImageService
+from kit.ai.model_builder import Service
+from kit.ai.pool import AIServicePool
+from kit.ai.rerank_service import RerankHit, RerankService
+
+__all__ = [
+    "FailoverChatService",
+    "FailoverEmbeddingService",
+    "FailoverImageService",
+    "FailoverRerankService",
+]
+
+T = TypeVar("T")
+
+# 批量向量化默认并发（与 EmbeddingService 保持一致）
+_EMBED_CONCURRENCY = 4
+
+
+class _FailoverBase:
+    """链式调用基类：候选迭代 + 失败冷却 + 聚合异常。"""
+
+    def __init__(self, pool: AIServicePool, model_type: str) -> None:
+        self._pool: AIServicePool = pool
+        self._model_type: str = model_type
+
+    async def _attempt(
+        self,
+        label: str,
+        run: Callable[[Service], Awaitable[T]],
+    ) -> T:
+        """按候选顺序执行 ``run``；单档案失败冷却后换下一个，全失败聚合抛出。"""
+        tried: list[str] = []
+        last_exc: AIError | None = None
+        for name, service in self._pool.iter_services(self._model_type):
+            tried.append(name)
+            try:
+                result = await run(service)
+            except AIError as exc:
+                last_exc = exc
+                self._pool.mark_failure(self._model_type, name, exc)
+                continue
+            self._pool.mark_success(self._model_type, name)
+            return result
+        raise AIRequestError(
+            f"{label}失败：已尝试 {len(tried)} 个档案"
+            f"（{', '.join(tried) or '无可用候选'}），最后错误: {last_exc}"
+        ) from last_exc
+
+
+class FailoverChatService(_FailoverBase):
+    """对话服务链式调用（ask / complete / stream）。"""
+
+    def __init__(self, pool: AIServicePool) -> None:
+        super().__init__(pool, "chat")
+
+    async def ask(
+        self,
+        messages: list[dict[str, str]],
+        *,
+        temperature: float | None = None,
+        max_tokens: int | None = None,
+    ) -> str:
+        """多轮对话，返回模型回答文本（链式故障切换）。"""
+
+        async def _run(service: Service) -> str:
+            return await cast(ChatService, service).ask(
+                messages, temperature=temperature, max_tokens=max_tokens
+            )
+
+        return await self._attempt("对话", _run)
+
+    async def complete(
+        self,
+        prompt: str,
+        *,
+        temperature: float | None = None,
+        max_tokens: int | None = None,
+    ) -> str:
+        """单轮对话，返回模型回答文本（链式故障切换）。"""
+
+        async def _run(service: Service) -> str:
+            return await cast(ChatService, service).complete(
+                prompt, temperature=temperature, max_tokens=max_tokens
+            )
+
+        return await self._attempt("单轮对话", _run)
+
+    async def stream(
+        self,
+        messages: list[dict[str, str]],
+        *,
+        temperature: float | None = None,
+        max_tokens: int | None = None,
+    ) -> AsyncIterator[str]:
+        """流式对话：首帧前失败自动换档案重试；首帧后失败直接抛出。"""
+        tried: list[str] = []
+        last_exc: AIError | None = None
+        for name, service in self._pool.iter_services(self._model_type):
+            tried.append(name)
+            agen = cast(
+                AsyncGenerator[str, None],
+                cast(ChatService, service).stream(
+                    messages, temperature=temperature, max_tokens=max_tokens
+                ),
+            )
+            try:
+                first = await agen.__anext__()
+            except StopAsyncIteration:
+                # 空流（首帧即结束）视为正常完成
+                await agen.aclose()
+                return
+            except AIError as exc:
+                last_exc = exc
+                self._pool.mark_failure(self._model_type, name, exc)
+                await agen.aclose()
+                continue
+            try:
+                self._pool.mark_success(self._model_type, name)
+                yield first
+                async for delta in agen:
+                    yield delta
+            finally:
+                await agen.aclose()
+            return
+        raise AIRequestError(
+            f"流式对话失败：已尝试 {len(tried)} 个档案"
+            f"（{', '.join(tried) or '无可用候选'}），最后错误: {last_exc}"
+        ) from last_exc
+
+
+class FailoverEmbeddingService(_FailoverBase):
+    """向量化服务链式调用（embed_text / embed_texts）。"""
+
+    def __init__(self, pool: AIServicePool) -> None:
+        super().__init__(pool, "embedding")
+
+    async def embed_text(self, text: str) -> list[float]:
+        """单文本向量化（链式故障切换）。"""
+        vectors = await self.embed_texts([text])
+        return vectors[0]
+
+    async def embed_texts(
+        self,
+        texts: list[str],
+        *,
+        batch_size: int | None = None,
+        concurrency: int = _EMBED_CONCURRENCY,
+    ) -> list[list[float]]:
+        """批量向量化（链式故障切换；整批失败才切换档案）。"""
+        if not texts:
+            return []
+
+        async def _run(service: Service) -> list[list[float]]:
+            return await cast(EmbeddingService, service).embed_texts(
+                texts, batch_size=batch_size, concurrency=concurrency
+            )
+
+        return await self._attempt("向量化", _run)
+
+
+class FailoverRerankService(_FailoverBase):
+    """重排序服务链式调用（rerank）。"""
+
+    def __init__(self, pool: AIServicePool) -> None:
+        super().__init__(pool, "rerank")
+
+    async def rerank(
+        self,
+        query: str,
+        documents: list[str],
+        *,
+        top_n: int | None = None,
+    ) -> list[RerankHit]:
+        """候选文档重排序（链式故障切换）。"""
+        if not documents:
+            return []
+
+        async def _run(service: Service) -> list[RerankHit]:
+            return await cast(RerankService, service).rerank(
+                query, documents, top_n=top_n
+            )
+
+        return await self._attempt("重排序", _run)
+
+
+class FailoverImageService(_FailoverBase):
+    """图像生成服务链式调用（generate）。"""
+
+    def __init__(self, pool: AIServicePool) -> None:
+        super().__init__(pool, "image")
+
+    async def generate(
+        self,
+        prompt: str,
+        *,
+        size: str | None = None,
+        n: int = 1,
+    ) -> str:
+        """文生图，返回图片 URL（链式故障切换）。"""
+
+        async def _run(service: Service) -> str:
+            return await cast(ImageService, service).generate(
+                prompt, size=size, n=n
+            )
+
+        return await self._attempt("图像生成", _run)

kit/ai/image_service.py

@@ -0,0 +1,99 @@
+"""
+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 override
+
+from kit.ai.base_service import BaseAIService, safe_get
+from kit.ai.exceptions import AIResponseError
+from kit.ai.types import JSONObject
+
+__all__ = ["ImageRequest", "ImageService"]
+
+
+@dataclass(frozen=True)
+class ImageRequest:
+    """文生图请求模型。
+
+    ``return_base64`` 对应 Agnes 官方参数 ``return_base64``，为 ``True``
+    时服务端以 ``data[0].b64_json`` 返回 Base64 数据而非 URL。
+    ``ratio`` 与档位式 ``size`` 配合的宽高比，官方支持
+    ``1:1/3:4/4:3/16:9/9:16/2:3/3:2/21:9``，默认 ``1:1``。
+    """
+
+    prompt: str
+    size: str | None = None
+    n: int = 1
+    return_base64: bool = False
+    ratio: str | None = None
+
+
+class ImageService(BaseAIService[ImageRequest, str]):
+    """图像生成服务（全异步）。
+
+    ``n > 1`` 时仅返回第一张图片，调用方如需多图请自行扩展。
+    返回 ``str`` 语义：默认返回图片 URL；``return_base64=True`` 时
+    返回 Base64 数据，调用方按返回值是否以 ``http`` 开头区分形式。
+    """
+
+    @override
+    def _build_payload(self, request: ImageRequest) -> JSONObject:
+        payload: JSONObject = {
+            "model": self._profile.model,
+            "prompt": request.prompt,
+            "n": request.n,
+        }
+        if request.size is not None:
+            payload["size"] = request.size
+        if request.return_base64:
+            payload["return_base64"] = True
+        if request.ratio is not None:
+            payload["ratio"] = request.ratio
+        return payload
+
+    @override
+    def _parse_response(self, data: JSONObject) -> str:
+        # 优先 URL 形式（data[0].url）；Base64 形式下 url 为 null 或整个键缺失
+        # （如 sensenova 响应仅含 b64_json），回退取 data[0].b64_json，
+        # 两者均无效才判定响应异常。
+        first = safe_get(data, "data", 0)
+        if isinstance(first, dict):
+            url = first.get("url")
+            if isinstance(url, str):
+                return url
+            b64 = first.get("b64_json")
+            if isinstance(b64, str):
+                return b64
+        raise AIResponseError(
+            f"图像生成响应缺少 url/b64_json: {str(data)[:200]}"
+        )
+
+    async def generate(
+        self,
+        prompt: str,
+        *,
+        size: str | None = None,
+        n: int = 1,
+        return_base64: bool = False,
+        ratio: str | None = None,
+    ) -> str:
+        request = ImageRequest(
+            prompt, size=size, n=n, return_base64=return_base64, ratio=ratio
+        )
+        return await self._call(request)

kit/ai/model_builder.py

@@ -0,0 +1,119 @@
+"""
+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 kit.ai.chat_service import ChatService
+from kit.ai.embedding_service import EmbeddingService
+from kit.ai.image_service import ImageService
+from kit.ai.profiles import ModelProfile, get_model_registry, with_api_key
+from kit.ai.rerank_service import RerankService
+
+__all__ = ["ModelBuilder"]
+
+# 支持的模型类型
+_MODEL_TYPES: tuple[str, ...] = ("chat", "embedding", "rerank", "image")
+
+Service = ChatService | EmbeddingService | RerankService | ImageService
+
+
+class ModelBuilder:
+    """根据模型类型构建对应的 AI 服务实例。"""
+
+    @staticmethod
+    def _resolve_profile(
+        model_type: str,
+        model: str | None = None,
+        *,
+        api_key: str | None = None,
+    ) -> ModelProfile:
+        registry = get_model_registry()
+        profile = registry.get(model_type, model)
+        if api_key is not None:
+            profile = with_api_key(profile, api_key)
+        return profile
+
+    @staticmethod
+    def build(
+        model_type: str,
+        model: str | None = None,
+        *,
+        api_key: str | None = None,
+    ) -> Service:
+        """构建指定的 AI 服务。
+
+        Args:
+            model_type: 模型类型（chat / embedding / rerank / image）
+            model: 档案名，None 时回退到该类型默认档案
+            api_key: 可选，覆盖档案中的密钥
+
+        Raises:
+            ValueError: 未知模型类型。
+        """
+        match model_type:
+            case "chat":
+                profile = ModelBuilder._resolve_profile(
+                    model_type, model, api_key=api_key
+                )
+                return ChatService(profile)
+            case "embedding":
+                profile = ModelBuilder._resolve_profile(
+                    model_type, model, api_key=api_key
+                )
+                return EmbeddingService(profile)
+            case "rerank":
+                profile = ModelBuilder._resolve_profile(
+                    model_type, model, api_key=api_key
+                )
+                return RerankService(profile)
+            case "image":
+                profile = ModelBuilder._resolve_profile(
+                    model_type, model, api_key=api_key
+                )
+                return ImageService(profile)
+            case _:
+                raise ValueError(
+                    f"未知模型类型: {model_type}，支持: {', '.join(_MODEL_TYPES)}"
+                )
+
+    @classmethod
+    def build_chat(
+        cls, model: str | None = None, *, api_key: str | None = None
+    ) -> ChatService:
+        profile = cls._resolve_profile("chat", model, api_key=api_key)
+        return ChatService(profile)
+
+    @classmethod
+    def build_embedding(
+        cls, model: str | None = None, *, api_key: str | None = None
+    ) -> EmbeddingService:
+        profile = cls._resolve_profile("embedding", model, api_key=api_key)
+        return EmbeddingService(profile)
+
+    @classmethod
+    def build_rerank(
+        cls, model: str | None = None, *, api_key: str | None = None
+    ) -> RerankService:
+        profile = cls._resolve_profile("rerank", model, api_key=api_key)
+        return RerankService(profile)
+
+    @classmethod
+    def build_image(
+        cls, model: str | None = None, *, api_key: str | None = None
+    ) -> ImageService:
+        profile = cls._resolve_profile("image", model, api_key=api_key)
+        return ImageService(profile)

kit/ai/model_client.py

@@ -0,0 +1,214 @@
+"""
+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
+"""
+
+import asyncio
+import json
+import random
+from collections.abc import AsyncIterator
+from typing import Final
+
+import httpx
+
+from kit.ai.exceptions import (
+    AIRequestError,
+    AIResponseError,
+    AITimeoutError,
+)
+from kit.ai.profiles import ModelProfile
+from kit.ai.types import JSONObject
+
+__all__ = ["ModelClient"]
+
+# post_json 对超时/5xx 的重试次数（指数退避）
+_DEFAULT_RETRIES = 2
+
+# list_models 轻探测独立超时（秒）：仅验证密钥有效性，与业务请求超时无关
+_LIST_MODELS_TIMEOUT: Final[float] = 10.0
+
+
+def _retryable(exc: BaseException) -> bool:
+    """判断异常是否可重试：超时，或服务端 5xx（按 AI 异常族状态码判定）。"""
+    if isinstance(exc, AITimeoutError):
+        return True
+    return (
+        isinstance(exc, AIRequestError)
+        and exc.status_code is not None
+        and exc.status_code >= 500
+    )
+
+
+class ModelClient:
+    """模型 HTTP 客户端（全异步，单客户端）。
+
+    普通请求与流式请求共用同一个 ``httpx.AsyncClient``；流式请求在
+    调用处通过 ``timeout=None`` 局部覆盖超时，避免双客户端维护成本。
+    ``post_json`` 对超时/5xx 做指数退避重试；``stream_json`` 不重试
+    （半开流连接重放语义复杂）。
+    """
+
+    def __init__(self, profile: ModelProfile) -> None:
+        self._profile: ModelProfile = profile
+        self._client: httpx.AsyncClient = httpx.AsyncClient(
+            base_url=profile.base_url,
+            timeout=profile.timeout,
+        )
+
+    @property
+    def path(self) -> str:
+        return self._profile.endpoint
+
+    @property
+    def model(self) -> str:
+        return self._profile.model
+
+    @staticmethod
+    def _headers(profile: ModelProfile) -> dict[str, str]:
+        return {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {profile.api_key}",
+        }
+
+    @staticmethod
+    async def _parse_json(data: str) -> JSONObject:
+        try:
+            parsed = json.loads(data)
+        except (json.JSONDecodeError, TypeError) as exc:
+            raise AIResponseError(f"响应不是合法 JSON: {exc}") from exc
+        if not isinstance(parsed, dict):
+            raise AIResponseError(f"响应顶层必须是 JSON 对象: {data[:200]}")
+        return parsed
+
+    @staticmethod
+    async def _parse_sse_line(line: str) -> JSONObject | None:
+        line = line.strip()
+        if not line or not line.startswith("data:"):
+            return None
+        payload = line[len("data:") :].strip()
+        if payload == "[DONE]":
+            return None
+        if payload == "[ERROR]":
+            raise AIResponseError("流式响应返回 [ERROR]")
+        return await ModelClient._parse_json(payload)
+
+    async def _post_once(self, payload: JSONObject) -> JSONObject:
+        """执行单次请求并解析响应体（不重试）。"""
+        try:
+            resp = await self._client.post(
+                self.path,
+                headers=self._headers(self._profile),
+                json=payload,
+            )
+            resp.raise_for_status()
+            return await self._parse_json(resp.text)
+        except httpx.TimeoutException as exc:
+            raise AITimeoutError(f"请求超时: {exc}") from exc
+        except httpx.HTTPStatusError as exc:
+            raise AIRequestError(
+                f"HTTP 错误 {exc.response.status_code}: {exc.response.text[:200]}",
+                status_code=exc.response.status_code,
+            ) from exc
+        except httpx.HTTPError as exc:
+            raise AIRequestError(f"请求失败: {exc}") from exc
+
+    async def post_json(
+        self,
+        payload: JSONObject,
+        *,
+        retries: int = _DEFAULT_RETRIES,
+    ) -> JSONObject:
+        """发送普通 JSON 请求并解析响应体。
+
+        对超时/5xx 做指数退避重试，累计 ``retries`` 次后抛出最后一次异常。
+        """
+        for attempt in range(retries + 1):
+            try:
+                return await self._post_once(payload)
+            except (AITimeoutError, AIRequestError) as exc:
+                if attempt >= retries or not _retryable(exc):
+                    raise
+                delay = min(2**attempt, 8) + random.uniform(0, 0.5)
+                await asyncio.sleep(delay)
+        # 不可达，供类型系统收口
+        raise AIRequestError("重试后仍未成功")
+
+    async def stream_json(
+        self, payload: JSONObject
+    ) -> AsyncIterator[JSONObject | None]:
+        """发送流式请求，逐帧产出已解析的 SSE 数据体。
+
+        空行与 ``[DONE]`` 标记会被跳过（产出 None 或不产出）；
+        ``[ERROR]`` 标记会抛出 ``AIResponseError``。不做重试。
+        """
+        try:
+            async with self._client.stream(
+                "POST",
+                self.path,
+                headers=self._headers(self._profile),
+                json=payload,
+                timeout=None,
+            ) as resp:
+                resp.raise_for_status()
+                async for line in resp.aiter_lines():
+                    parsed = await self._parse_sse_line(line)
+                    if parsed is not None:
+                        yield parsed
+        except httpx.TimeoutException as exc:
+            raise AITimeoutError(f"流式请求超时: {exc}") from exc
+        except httpx.HTTPStatusError as exc:
+            raise AIRequestError(
+                f"流式 HTTP 错误 {exc.response.status_code}: {exc.response.text[:200]}",
+                status_code=exc.response.status_code,
+            ) from exc
+        except httpx.HTTPError as exc:
+            raise AIRequestError(f"流式请求失败: {exc}") from exc
+
+    async def list_models(self) -> JSONObject:
+        """GET ``{base_url}/models`` 轻探测：仅验证密钥有效性（候选链预热用）。
+
+        独立 10 秒超时（不影响业务请求配置）；使用绝对 URL 直连——
+        ``base_url`` 已含版本前缀，若用相对路径会被 httpx 按根路径拼接。
+
+        Raises:
+            AIRequestError: 非 2xx（``status_code`` 携带 HTTP 状态码）
+                或连接失败。
+            AITimeoutError: 探测超时。
+            AIResponseError: 响应不是合法 JSON 对象。
+        """
+        url = self._profile.base_url.rstrip("/") + "/models"
+        try:
+            resp = await self._client.get(
+                url,
+                headers=self._headers(self._profile),
+                timeout=_LIST_MODELS_TIMEOUT,
+            )
+            resp.raise_for_status()
+            return await self._parse_json(resp.text)
+        except httpx.TimeoutException as exc:
+            raise AITimeoutError(f"模型探测超时: {exc}") from exc
+        except httpx.HTTPStatusError as exc:
+            raise AIRequestError(
+                f"模型探测 HTTP 错误 {exc.response.status_code}: "
+                f"{exc.response.text[:200]}",
+                status_code=exc.response.status_code,
+            ) from exc
+        except httpx.HTTPError as exc:
+            raise AIRequestError(f"模型探测请求失败: {exc}") from exc
+
+    async def close(self) -> None:
+        await self._client.aclose()