112 lines
4.0 KiB
Python
112 lines
4.0 KiB
Python
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)
|