123 lines
4.5 KiB
Python
123 lines
4.5 KiB
Python
from collections.abc import Mapping
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from app.config import Settings
|
|
|
|
|
|
@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],
|
|
tests: Mapping[str, Mapping[str, Any]] | None = None,
|
|
) -> None:
|
|
self.settings = settings
|
|
self.tests = tests
|
|
|
|
def _value(self, key: str, default: Any = "") -> Any:
|
|
if isinstance(self.settings, Mapping):
|
|
return self.settings.get(key, default)
|
|
return getattr(self.settings, key, default)
|
|
|
|
@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 [
|
|
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:
|
|
return str(self.choose("image").get("provider", "custom"))
|