初始化提交

backend/app/ai_provider/routing/selector.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+from typing import Any
+
+from app.ai_provider.config.models import ModelConfig
+from app.ai_provider.core.registry import ModelRegistry
+from app.ai_provider.routing.multimodal import MultiModalDecision
+from app.ai_provider.types.errors import AIProviderError
+
+_SELECTOR_CAPABILITY_MAP: dict[MultiModalDecision, str] = {
+    MultiModalDecision.VISION: "vision",
+    MultiModalDecision.THINKING: "thinking",
+    MultiModalDecision.TOOL_USE: "tool_use",
+    MultiModalDecision.STRUCTURED_OUTPUT: "structured_output",
+    MultiModalDecision.PLAIN_TEXT: "chat",
+}
+
+
+class ModelSelector:
+    """为给定上下文选择最佳模型。
+
+    用法::
+
+        selector = ModelSelector(registry)
+        model = selector.select({"capabilities": {"vision", "thinking"}})
+        model = selector.select({"preferred": "fast-model"})
+        model = selector.select_from_decisions({MultiModalDecision.VISION})
+    """
+
+    def __init__(self, registry: ModelRegistry) -> None:
+        self._registry = registry
+
+    def select(self, context: dict[str, Any]) -> ModelConfig:
+        """根据给定的上下文选择模型。
+
+        上下文键:
+
+        * ``capabilities`` (``set[str]``, 可选) — 必需的能力。
+        * ``preferred`` (``str``, 可选) — 偏好的模型名或别名。
+        * ``avoid`` (``set[str]``, 可选) — 要排除的模型名。
+
+        如果给了偏好模型但不能满足所有必需的能力，此方法会回退到
+        能满足要求的模型。
+        """
+        capabilities: set[str] = set(context.get("capabilities", set()))
+        preferred: str | None = context.get("preferred")
+        avoid: set[str] = set(context.get("avoid", set()))
+
+        # If preferred is given, resolve it and check capabilities
+        if preferred is not None:
+            model = self._registry.resolve(preferred)
+            if (
+                not capabilities or self._model_satisfies_all(model, capabilities)
+            ) and model.name not in avoid:
+                return model
+
+        # Find models that satisfy all required capabilities
+        if capabilities:
+            return self._select_by_capabilities(capabilities, avoid)
+
+        # Fall back to default
+        return self._registry.resolve()
+
+    def select_from_decisions(
+        self,
+        decisions: set[MultiModalDecision],
+    ) -> ModelConfig:
+        """基于 :class:`MultiModalDecision` 值选择模型。
+
+        将决策翻译为能力字符串并委托给 :meth:`select`。
+        """
+        capabilities: set[str] = set()
+        for d in decisions:
+            cap = _SELECTOR_CAPABILITY_MAP.get(d)
+            if cap and cap != "chat":
+                capabilities.add(cap)
+        return self.select({"capabilities": capabilities})
+
+    # ------------------------------------------------------------------
+
+    def _select_by_capabilities(
+        self,
+        capabilities: set[str],
+        avoid: set[str],
+    ) -> ModelConfig:
+        """找到支持所有必需能力的模型。"""
+        # Narrow down by each capability (ModelConfig is unhashable, use lists)
+        candidates: list[ModelConfig] = list(self._registry.list_all())
+
+        for cap in capabilities:
+            cap_models = list(self._registry.list_by_capability(cap))
+            cap_names = {m.name for m in cap_models}
+            candidates = [m for m in candidates if m.name in cap_names]
+
+        # Remove avoided models
+        candidates = [m for m in candidates if m.name not in avoid]
+
+        if not candidates:
+            raise AIProviderError(
+                message=f"No active model supports all capabilities: {capabilities}",
+                code="CAPABILITY_NOT_FOUND",
+                status_code=404,
+            )
+
+        # Prefer higher-weight models (higher = better)
+        return max(candidates, key=lambda m: m.weight)
+
+    @staticmethod
+    def _model_satisfies_all(model: ModelConfig, capabilities: set[str]) -> bool:
+        caps = model.resolved_capabilities()
+        return all(caps.get(cap, False) for cap in capabilities)