sdpy-kit 0.1.0__py3-none-any.whl

kit/ai/runner.py

@@ -0,0 +1,580 @@
+"""
+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 __future__ import annotations
+
+# date: 2026-08-24
+"""kit/ai 同步门面（共享后台事件循环 + 常驻链式服务单例）
+
+``kit.ai`` 各服务（ChatService / EmbeddingService / RerankService / ImageService）
+均为全异步，且实例持有绑定事件循环的 ``httpx.AsyncClient``；而业务层（service）
+多数调用点运行在同步上下文（如 ``asyncio.to_thread`` 线程、被 FastAPI 路由直调的
+同步函数）。
+
+本模块用 Python 3.13 标准库实现统一适配：
+
+- 一个由 ``asyncio.Runner`` 驱动的**共享后台事件循环线程**（daemon），常驻复用，
+  避免每次调用反复创建线程池与事件循环。
+- 通过 ``asyncio.run_coroutine_threadsafe`` 把协程从**同步上下文**（含线程、
+  ``asyncio.to_thread`` 线程）投递到后台 loop 并同步阻塞取结果——无需判断
+  调用线程是否有事件循环，也规避了 ``asyncio.run()`` 在已运行 loop 内被禁的坑。
+- 按模型类型**缓存常驻 Failover 链式服务单例**：内部经 ``AIServicePool``
+  按候选链顺序调用（优先级链 + 失败冷却 + 自动切换），并复用底层
+  ``httpx.AsyncClient`` 连接池。所有 I/O 都发生在同一后台 loop。
+
+用法::
+
+    from kit.ai.runner import AI, shutdown
+    text = AI.chat_complete(prompt)      # 同步
+    vectors = AI.embed_texts(chunks)     # 同步
+    hits = AI.rerank(query, docs, top_n) # 同步
+    url = AI.generate_image(prompt)      # 同步
+    AI.use("chat", "glm")                # 切换 chat 当前模型档案（长期生效）
+    print(AI.available("chat"))          # 查看各档案状态快照
+    # 应用退出时
+    shutdown()
+
+注意：
+- 本门面仅面向**同步上下文**。若调用点本身是 async 函数（事件循环内），
+  应直接 ``await`` 原生 ``ChatService.ask/stream`` 等异步接口，
+  而非经本门面阻塞事件循环线程。
+- ``chat_stream`` 提供同步流式迭代（每次取下一帧），适合同步线程内逐 token 消费。
+- ``AI.use`` / ``AI.available`` 为进程内内存态（不持久化），进程重启后
+  回到 env（``AI_{TYPE}_PROFILES``）/ JSON 声明顺序。
+"""
+
+import asyncio
+import concurrent.futures
+import queue
+import threading
+from collections.abc import Callable, Coroutine, Iterator
+from typing import Any, Final, TypeVar, cast
+
+from kit.ai.exceptions import AIConfigError
+from kit.ai.failover import (
+    FailoverChatService,
+    FailoverEmbeddingService,
+    FailoverImageService,
+    FailoverRerankService,
+)
+from kit.ai.pool import AIServicePool
+from kit.ai.profiles import ProfileStatus, get_model_registry
+from kit.ai.rerank_service import RerankHit
+from kit.config.log_config import get_default_logger
+
+T = TypeVar("T")
+
+logger = get_default_logger(__name__)
+
+# 后台常驻服务类型键
+_CHAT = "chat"
+_EMBEDDING = "embedding"
+_RERANK = "rerank"
+_IMAGE = "image"
+
+# 支持的模型类型集合（use/available 校验与错误提示用）
+_SERVICE_TYPES: Final[frozenset[str]] = frozenset(
+    {_CHAT, _EMBEDDING, _RERANK, _IMAGE}
+)
+
+# 常驻链式服务联合类型（按模型类型缓存于后台线程）
+FailoverService = (
+    FailoverChatService
+    | FailoverEmbeddingService
+    | FailoverRerankService
+    | FailoverImageService
+)
+
+# 同步门面单次调用的默认等待上限（秒），兜底防进程退出/网络挂起；
+# 各类型实际等待上限由 _facade_timeout 按档案 timeout 自适应（+ 余量）。
+_RUN_TIMEOUT = 120.0
+
+# 门面等待上限相对模型侧 timeout 的余量（秒）
+_RUN_TIMEOUT_MARGIN: Final[float] = 10.0
+
+# ``chat_stream`` 流结束哨兵（放入线程安全队列）
+_SENTINEL = object()
+
+
+class _BackgroundLoop:
+    """共享后台事件循环：单线程常驻，从任意线程安全投递协程。
+
+    Python 3.13 实现：``asyncio.Runner`` 持有并驱动一个专属事件循环，
+    ``run_coroutine_threadsafe`` 负责线程间投递与结果同步。
+    """
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._started = threading.Event()
+        self._thread: threading.Thread | None = None
+        self._runner: asyncio.Runner | None = None
+        self._loop: asyncio.AbstractEventLoop | None = None
+        # 后台 loop 内的 keep_alive 停止信号（线程间经 call_soon_threadsafe 操作）
+        self._keep_alive_event: asyncio.Event | None = None
+        # 仅在后台线程内读写，主线程只经 shutdown() 间接操作
+        self._services: dict[str, FailoverService] = {}
+        # AI 候选池（惰性创建；业务线程 AI.use/AI.available 与 loop 线程共享）
+        self._pool: AIServicePool | None = None
+        # 预热探测任务（进程内仅投递一次，持引用防 GC）
+        self._probe_task: asyncio.Task[None] | None = None
+
+    # ---------- 生命周期 ----------
+
+    def _ensure_started(self) -> asyncio.AbstractEventLoop:
+        loop = self._loop
+        if loop is not None and not loop.is_closed():
+            return loop
+        with self._lock:
+            loop = self._loop
+            if loop is not None and not loop.is_closed():
+                return loop
+            self._started = threading.Event()
+            self._thread = threading.Thread(
+                target=self._worker, name="kit-ai-runner", daemon=True
+            )
+            self._thread.start()
+        # 在 lock 外等待，避免与 worker 内获取 lock 设置 _loop 相互死锁
+        self._started.wait(timeout=5)
+        loop = self._loop
+        if loop is None:
+            raise RuntimeError("kit/ai 后台事件循环启动失败")
+        return loop
+
+    def _worker(self) -> None:
+        # Runner 需在专属线程内创建（loop 绑定创建线程）
+        runner = asyncio.Runner()
+        loop = runner.get_loop()
+        with self._lock:
+            self._runner = runner
+            self._loop = loop
+            self._keep_alive_event = asyncio.Event()
+        self._started.set()
+        try:
+            # keep_alive 阻塞等待停止信号，让 runner.run 持续处理投递的协程
+            runner.run(self._keep_alive())
+        except asyncio.CancelledError:
+            logger.debug("kit/ai 后台事件循环被取消")
+        finally:
+            self._close_services()
+            with self._lock:
+                runner.close()
+                self._runner = None
+                self._loop = None
+                self._keep_alive_event = None
+
+    async def _keep_alive(self) -> None:
+        event = self._keep_alive_event
+        if event is not None:
+            await event.wait()
+
+    def _close_services(self) -> None:
+        """在后台线程内关闭候选池中全部常驻服务（关闭底层 AsyncClient）。
+
+        Failover 链式服务本身不持连接，连接在池内各档案的 service 上，
+        统一经 ``pool.close_all()`` 释放。
+        """
+        loop = self._loop
+        if loop is None or loop.is_closed():
+            return
+        self._services = {}
+        pool = self._pool
+        if pool is None:
+            return
+        try:
+            loop.run_until_complete(pool.close_all())
+        except Exception as exc:  # noqa: BLE001 - 关闭阶段失败不影响整体退出
+            logger.warning(f"关闭 AI 候选池失败: {exc}")
+
+    def shutdown(self) -> None:
+        """停止后台事件循环并释放常驻服务。可在任意线程调用，幂等。"""
+        with self._lock:
+            loop = self._loop
+            thread = self._thread
+            event = self._keep_alive_event
+        # 触发 keep_alive 结束，让 runner.run 干净返回，避免 pending task 泄漏
+        if loop is not None and loop.is_running() and event is not None:
+            loop.call_soon_threadsafe(event.set)
+        if thread is not None and thread is not threading.current_thread():
+            thread.join(timeout=5)
+
+    # ---------- 候选池与当前项切换 ----------
+
+    def _get_pool(self) -> AIServicePool:
+        """取（或首次创建）候选池：锁内单例，业务线程与 loop 线程共享。"""
+        with self._lock:
+            if self._pool is None:
+                self._pool = AIServicePool()
+            return self._pool
+
+    def _schedule_probe(self) -> None:
+        """进程内首次取服务时投递候选链预热探测（fire-and-forget，不阻塞）。
+
+        ``probe_all`` 逐类型独立容错且自身不抛异常；持引用防任务被 GC。
+        """
+        if self._probe_task is not None:
+            return
+        pool = self._get_pool()
+        self._probe_task = asyncio.create_task(pool.probe_all())
+
+    def use(self, model_type: str, profile: str) -> None:
+        """切换该类型长期当前档案：指定项成为固定链头并清除冷却。
+
+        同步方法，可从任意业务线程直接调用（池内部以锁保护共享状态）。
+
+        Raises:
+            AIConfigError: 类型不存在、档案不存在或不可用
+                （未配置 / 密钥无效 / 不可达 / 因约束被排除）。
+        """
+        self._get_pool().prefer(model_type, profile)
+
+    def available(self, model_type: str) -> list[ProfileStatus]:
+        """返回该类型全部档案状态快照（含未配置者），链头标记 current。"""
+        return self._get_pool().available(model_type)
+
+    # ---------- 投递 ----------
+
+    def run(
+        self,
+        coro_factory: Callable[[], Coroutine[Any, Any, T]],
+        *,
+        timeout: float = _RUN_TIMEOUT,
+    ) -> T:
+        """同步执行协程工厂：投递到后台 loop 并阻塞取结果。
+
+        因使用独立后台 loop，从同步代码、线程内均可安全调用
+        （``run_coroutine_threadsafe`` 线程安全）。
+
+        Args:
+            coro_factory: 协程工厂（每次调用产出一个新协程）
+            timeout: 同步等待上限（秒）；应不低于协程内部模型侧超时
+        """
+        loop = self._ensure_started()
+        if threading.current_thread() is self._thread:
+            # 避免在后台线程内二次同步等待造成死锁
+            raise RuntimeError("kit/ai 同步门面不支持在后台线程内再次同步调用")
+        coro = coro_factory()
+        future: concurrent.futures.Future[T] = asyncio.run_coroutine_threadsafe(
+            coro, loop
+        )
+        try:
+            return future.result(timeout=timeout)
+        except concurrent.futures.TimeoutError:
+            # 后台 loop 内的协程仍在运行，取消它避免泄漏
+            future.cancel()
+            raise TimeoutError(f"kit/ai 同步调用超过 {timeout:g}s 未返回") from None
+
+    def submit_stream(
+        self,
+        messages: list[dict[str, str]],
+        q: queue.Queue[object],
+    ) -> concurrent.futures.Future[None]:
+        """投递流式收集任务：后台拉取链式 ``stream`` 帧写入线程安全队列。
+
+        返回的 future 可取消，用于同步侧中断时停止后台拉取。
+        """
+        loop = self._ensure_started()
+
+        async def _producer() -> None:
+            service: FailoverChatService = cast(
+                FailoverChatService, await self._get_service(_CHAT)
+            )
+            try:
+                async for delta in service.stream(messages):
+                    q.put(delta)
+            except BaseException as exc:  # noqa: BLE001 - 异常经队列回传给同步侧
+                q.put(exc)
+            finally:
+                q.put(_SENTINEL)
+
+        return asyncio.run_coroutine_threadsafe(_producer(), loop)
+
+    async def _get_service(self, service_type: str) -> FailoverService:
+        """后台线程内取常驻链式服务，未创建则就地构建并缓存。
+
+        进程内首次构建任意类型时投递一次候选链预热探测（fire-and-forget）。
+        """
+        service = self._services.get(service_type)
+        if service is None:
+            self._schedule_probe()
+            pool = self._get_pool()
+            match service_type:
+                case "chat":
+                    service = FailoverChatService(pool)
+                case "embedding":
+                    service = FailoverEmbeddingService(pool)
+                case "rerank":
+                    service = FailoverRerankService(pool)
+                case "image":
+                    service = FailoverImageService(pool)
+                case _:
+                    raise ValueError(
+                        f"未知模型类型: {service_type}，"
+                        f"支持: {', '.join(sorted(_SERVICE_TYPES))}"
+                    )
+            self._services[service_type] = service
+            logger.debug(f"已构建并常驻 AI 链式服务：{service_type}")
+        return service
+
+
+_background = _BackgroundLoop()
+
+
+def run_sync[T](
+    coro_factory: Callable[[], Coroutine[Any, Any, T]],
+    *,
+    timeout: float = _RUN_TIMEOUT,
+) -> T:
+    """在任意上下文同步执行协程工厂（兼容旧接口，等价于后台投递）。"""
+    return _background.run(coro_factory, timeout=timeout)
+
+
+def _facade_timeout(model_type: str) -> float:
+    """门面同步等待上限：该类型档案 timeout 最大值 + 余量（秒）。
+
+    读静态档案（``type_profiles``，无需密钥），随 ``ai_models.json``
+    的超时配置自适应；类型未配置时回退默认 ``_RUN_TIMEOUT``。
+    """
+    try:
+        profiles = get_model_registry().type_profiles(model_type)
+    except AIConfigError:
+        return _RUN_TIMEOUT
+    if not profiles:
+        return _RUN_TIMEOUT
+    model_timeout = max(config.timeout for config in profiles.values())
+    return float(model_timeout) + _RUN_TIMEOUT_MARGIN
+
+
+async def _chat_ask(messages: list[dict[str, str]]) -> str:
+    service: FailoverChatService = cast(
+        FailoverChatService, await _background._get_service(_CHAT)  # noqa: SLF001
+    )
+    return await service.ask(messages)
+
+
+async def _chat_complete(prompt: str) -> str:
+    service: FailoverChatService = cast(
+        FailoverChatService, await _background._get_service(_CHAT)  # noqa: SLF001
+    )
+    return await service.complete(prompt)
+
+
+async def _embed_texts(texts: list[str]) -> list[list[float]]:
+    if not texts:
+        return []
+    service: FailoverEmbeddingService = cast(
+        FailoverEmbeddingService,
+        await _background._get_service(_EMBEDDING),  # noqa: SLF001
+    )
+    return await service.embed_texts(texts)
+
+
+async def _rerank(
+    query: str, documents: list[str], top_n: int | None
+) -> list[tuple[int, float]]:
+    if not documents:
+        return []
+    service: FailoverRerankService = cast(
+        FailoverRerankService,
+        await _background._get_service(_RERANK),  # noqa: SLF001
+    )
+    hits: list[RerankHit] = await service.rerank(query, documents, top_n=top_n)
+    return [(hit.index, hit.score) for hit in hits]
+
+
+async def _generate_image(prompt: str, *, size: str | None, n: int) -> str:
+    service: FailoverImageService = cast(
+        FailoverImageService,
+        await _background._get_service(_IMAGE),  # noqa: SLF001
+    )
+    return await service.generate(prompt, size=size, n=n)
+
+
+class AI:
+    """kit/ai 同步门面：在同步上下文调用异步 AI 服务（候选链故障切换）。"""
+
+    @staticmethod
+    def chat_complete(prompt: str) -> str:
+        """单轮对话（仅 user 消息），返回模型回答文本。"""
+        return _background.run(
+            lambda: _chat_complete(prompt), timeout=_facade_timeout(_CHAT)
+        )
+
+    @staticmethod
+    def chat_ask(messages: list[dict[str, str]]) -> str:
+        """多轮对话（自组装 messages），返回模型回答文本。"""
+        return _background.run(
+            lambda: _chat_ask(messages), timeout=_facade_timeout(_CHAT)
+        )
+
+    @staticmethod
+    def embed_text(text: str) -> list[float]:
+        """对单个文本向量化，返回向量。"""
+        return _background.run(
+            lambda: _embed_texts([text]), timeout=_facade_timeout(_EMBEDDING)
+        )[0]
+
+    @staticmethod
+    def embed_texts(texts: list[str]) -> list[list[float]]:
+        """对文本列表批量向量化，返回向量列表。"""
+        return _background.run(
+            lambda: _embed_texts(texts), timeout=_facade_timeout(_EMBEDDING)
+        )
+
+    @staticmethod
+    def rerank(
+        query: str, documents: list[str], top_n: int | None = None
+    ) -> list[tuple[int, float]]:
+        """对候选文档重排序，返回 [(原索引, 相关分)] 降序列表。"""
+        return _background.run(
+            lambda: _rerank(query, documents, top_n), timeout=_facade_timeout(_RERANK)
+        )
+
+    @staticmethod
+    def generate_image(prompt: str, *, size: str | None = None, n: int = 1) -> str:
+        """文生图，返回生成的图片 URL。"""
+        return _background.run(
+            lambda: _generate_image(prompt, size=size, n=n),
+            timeout=_facade_timeout(_IMAGE),
+        )
+
+    @staticmethod
+    def use(model_type: str, profile: str) -> None:
+        """切换某类型长期当前模型档案（进程内生效，重启回 env/json 序）。
+
+        指定档案成为固定链头并清除其冷却；再次 ``use`` 才会更换。
+
+        Args:
+            model_type: 模型类型（chat / embedding / rerank / image）
+            profile: 目标档案名（``ai_models.json`` 中该类型 profiles 的键）
+
+        Raises:
+            AIConfigError: 类型不存在、档案不存在或不可用
+                （未配置 / 密钥无效 / 不可达 / 因约束被排除）。
+        """
+        _background.use(model_type, profile)
+
+    @staticmethod
+    def available(model_type: str) -> list[ProfileStatus]:
+        """返回某类型全部档案状态快照（含未配置与被剔除者），链头标 current。
+
+        Args:
+            model_type: 模型类型（chat / embedding / rerank / image）
+
+        Returns:
+            ``ProfileStatus`` 列表（name/model/state/current/last_error/
+            last_ok_at），供业务展示与切换决策。
+
+        Raises:
+            AIConfigError: 类型不存在或注册表未加载。
+        """
+        return _background.available(model_type)
+
+    @staticmethod
+    def chat_stream(
+        messages: list[dict[str, str]],
+    ) -> Iterator[str]:
+        """多轮对话（流式），返回逐 token 迭代器（同步）。
+
+        在同步线程内逐帧消费，适合流式输出；内部在后台 loop 持续拉取
+        并放入线程安全队列。
+
+        推荐配合 ``with`` 使用，保证中途 ``break`` 也会及时取消后台任务::
+
+            with AI.chat_stream(messages) as stream:
+                for token in stream:
+                    yield token
+                    if done:
+                        break
+        """
+        return _StreamIterator(messages)
+
+
+class _StreamIterator:
+    """``AI.chat_stream`` 的同步流式迭代器。
+
+    通过 ``queue.Queue`` 桥接后台 loop：一个后台任务把链式 ``stream``
+    产出的帧逐条放入队列，同步侧阻塞 ``get`` 消费；停止信号由哨兵标记。
+
+    支持 ``with`` 上下文（推荐）：退出时主动取消后台任务；即便调用方
+    直接 ``for`` 中途 ``break`` 或迭代器被 GC，``close``/``__del__`` 也会
+    兜底取消后台 producer，避免任务残留。
+    """
+
+    def __init__(self, messages: list[dict[str, str]]) -> None:
+        self._messages: list[dict[str, str]] = messages
+        self._q: queue.Queue[object] = queue.Queue()
+        self._started = False
+        self._closed = False
+        self._task: concurrent.futures.Future[None] | None = None
+
+    def __iter__(self) -> Iterator[str]:
+        return self
+
+    def _ensure_started(self) -> None:
+        if not self._started:
+            self._started = True
+            self._task = _background.submit_stream(self._messages, self._q)
+
+    def _cancel_task(self) -> None:
+        task = self._task
+        if task is not None:
+            task.cancel()
+        self._task = None
+
+    def close(self) -> None:
+        """停止后台拉取并释放资源。幂等，可多次调用。"""
+        if self._closed:
+            return
+        self._closed = True
+        self._cancel_task()
+
+    def __del__(self) -> None:
+        # 兜底清理：for 中途 break / 迭代器被 GC 时取消后台 producer。
+        # 模块卸载期 _task 是独立 future 引用，cancel 对已 shutdown 的 loop 安全。
+        try:
+            self.close()
+        except Exception:  # noqa: BLE001 - GC 阶段不容许抛异常
+            pass
+
+    def __enter__(self) -> _StreamIterator:
+        return self
+
+    def __exit__(self, *exc_info: object) -> None:
+        self.close()
+
+    def __next__(self) -> str:
+        if self._closed:
+            raise StopIteration
+        self._ensure_started()
+        item = self._q.get()
+        if item is _SENTINEL:
+            self.close()
+            raise StopIteration
+        if isinstance(item, BaseException):
+            self.close()
+            raise item
+        return str(item)
+
+
+def shutdown() -> None:
+    """停止后台事件循环并释放常驻 AI 服务（应用退出前调用，幂等）。"""
+    _background.shutdown()
+
+
+__all__ = ["AI", "run_sync", "shutdown"]

kit/ai/types.py

@@ -0,0 +1,28 @@
+"""
+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
+"""
+
+"""kit/ai 公共类型别名（PEP 695 type 语句，递归 JSON 类型）"""
+
+# 递归 JSON 值类型：标量 | 数组 | 对象
+type JSONScalar = str | int | float | bool | None
+type JSONValue = JSONScalar | list[JSONValue] | dict[str, JSONValue]
+# JSON 对象（请求体 / 响应体）
+type JSONObject = dict[str, JSONValue]
+
+__all__ = ["JSONObject", "JSONScalar", "JSONValue"]