feat: add configurable AI model pool

2026-08-01 23:11:44 +08:00
parent 020b9596dc
commit 4430149ee4
13 changed files with 1660 additions and 273 deletions
@@ -1,43 +1,122 @@
-from dataclasses import dataclass
 from collections.abc import Mapping
+from dataclasses import dataclass
 from typing import Any

 from app.config import Settings


-def _configured(value: str) -> bool:
-    return bool(value) and not value.startswith("replace_with_")
-
-
@dataclass(frozen=True)
 class ProviderCapability:
    provider: str
    configured: bool
    strengths: tuple[str, ...]
+    model_id: str = ""
+    name: str = ""


 class ModelRouter:
-    def __init__(self, settings: Settings | Mapping[str, Any]) -> None:
+    def __init__(
+        self,
+        settings: Settings | Mapping[str, Any],
+        tests: Mapping[str, Mapping[str, Any]] | None = None,
+    ) -> None:
        self.settings = settings
+        self.tests = tests

-    def _value(self, key: str, default: str = "") -> str:
+    def _value(self, key: str, default: Any = "") -> Any:
        if isinstance(self.settings, Mapping):
-            return str(self.settings.get(key, default) or "")
-        return str(getattr(self.settings, key, default) or "")
+            return self.settings.get(key, default)
+        return getattr(self.settings, key, default)

-    def capabilities(self) -> list[ProviderCapability]:
+    @property
+    def pool(self) -> list[dict[str, Any]]:
+        raw_pool = self._value("model_pool", [])
+        return [item for item in raw_pool if isinstance(item, dict) and item.get("enabled", True)]
+
+    def _tested(self, item: Mapping[str, Any]) -> bool:
+        if self.tests is None:
+            return True
+        return bool(self.tests.get(f"model:{item.get('id')}", {}).get("ok"))
+
+    def candidates(self, capability: str) -> list[dict[str, Any]]:
        return [
-            ProviderCapability(
-                provider=self._value("ai_provider", "custom"),
-                configured=_configured(self._value("ai_api_key"))
-                and bool(self._value("ai_base_url"))
-                and bool(self._value("image_model")),
-                strengths=("统一模型网关", "空间理解", "图像生成与编辑"),
-            ),
+            item
+            for item in self.pool
+            if capability in item.get("capabilities", []) and self._tested(item)
        ]

+    def orchestrator(self) -> dict[str, Any]:
+        selected_id = str(self._value("orchestrator_model_id", ""))
+        selected = next(
+            (
+                item
+                for item in self.pool
+                if item.get("id") == selected_id
+                and item.get("category") == "language"
+                and self._tested(item)
+            ),
+            None,
+        )
+        if selected:
+            return selected
+        raise RuntimeError("总调度模型未配置或尚未通过测试。")
+
+    def choose(self, route: str, orchestrator_choice_id: str | None = None) -> dict[str, Any]:
+        if route == "spatial":
+            mode_key, selected_key, capability = (
+                "spatial_routing_mode",
+                "spatial_model_id",
+                "spatial_understanding",
+            )
+        elif route == "image":
+            mode_key, selected_key, capability = (
+                "image_routing_mode",
+                "image_model_id",
+                "image_generation",
+            )
+        else:
+            raise ValueError(f"未知模型路由：{route}")
+
+        candidates = self.candidates(capability)
+        if self._value(mode_key, "auto") == "manual":
+            selected_id = str(self._value(selected_key, ""))
+            selected = next((item for item in candidates if item.get("id") == selected_id), None)
+            if selected:
+                return selected
+            raise RuntimeError(f"手动选择的{route}模型不可用或尚未通过测试。")
+        if orchestrator_choice_id:
+            selected = next((item for item in candidates if item.get("id") == orchestrator_choice_id), None)
+            if selected:
+                return selected
+            raise RuntimeError("总调度返回的模型不在已测试候选列表中。")
+        if len(candidates) == 1:
+            return candidates[0]
+        if len(candidates) > 1:
+            raise RuntimeError("自动路由存在多个候选，需要总调度返回模型实例 ID。")
+        raise RuntimeError(f"没有通过测试且支持{capability}的模型。")
+
+    def capabilities(self) -> list[ProviderCapability]:
+        result: list[ProviderCapability] = []
+        for item in self.pool:
+            strengths: list[str] = []
+            if "orchestration" in item.get("capabilities", []):
+                strengths.append("总调度")
+            if "spatial_understanding" in item.get("capabilities", []):
+                strengths.append("空间理解")
+            if "image_generation" in item.get("capabilities", []):
+                strengths.append("图像生成")
+            if "image_editing" in item.get("capabilities", []):
+                strengths.append("图像编辑")
+            result.append(
+                ProviderCapability(
+                    provider=str(item.get("provider", "custom")),
+                    configured=self._tested(item),
+                    strengths=tuple(strengths),
+                    model_id=str(item.get("model_id", "")),
+                    name=str(item.get("name", "")),
+                )
+            )
+        return result
+
    def choose_image_provider(self, task: str) -> str:
-        configured = [item.provider for item in self.capabilities() if item.configured]
-        if configured:
-            return configured[0]
-        raise RuntimeError("No image provider is configured.")
+        return str(self.choose("image").get("provider", "custom"))