120 lines
4.9 KiB
Python
120 lines
4.9 KiB
Python
from __future__ import annotations
|
||
|
||
import os
|
||
from typing import Any
|
||
|
||
|
||
EXTERNAL_QUERY_PROVIDER_SPECS = {
|
||
"mcp": {
|
||
"id": "mcp",
|
||
"label": "MCP 外部查询 Adapter",
|
||
"kind": "mcp_adapter",
|
||
"summary": "优先通过已接入 MCP runtime 获取外部公开来源。",
|
||
},
|
||
"live_search": {
|
||
"id": "live_search",
|
||
"label": "Bounded Live Search Fallback",
|
||
"kind": "bounded_live_search",
|
||
"summary": "MCP 不可用时使用受控实时搜索 fallback，并继续经过 verifier 与来源过滤。",
|
||
},
|
||
}
|
||
|
||
|
||
def _env_enabled(name: str, default: bool = True) -> bool:
|
||
raw = str(os.getenv(name, "1" if default else "0") or "").strip().lower()
|
||
return raw not in {"0", "false", "no", "off", "disabled"}
|
||
|
||
|
||
class ExternalProviderRegistry:
|
||
def __init__(
|
||
self,
|
||
*,
|
||
adapters: list[Any] | None = None,
|
||
provider_statuses: list[dict[str, Any]] | None = None,
|
||
) -> None:
|
||
self._adapters = list(adapters or [])
|
||
self._provider_statuses = list(provider_statuses or [])
|
||
|
||
def iter_adapters(self) -> list[Any]:
|
||
return list(self._adapters)
|
||
|
||
def provider_statuses(self) -> list[dict[str, Any]]:
|
||
if self._provider_statuses:
|
||
return [dict(item) for item in self._provider_statuses if isinstance(item, dict)]
|
||
adapter_names = {
|
||
str(getattr(adapter, "adapter_name", "") or "").strip()
|
||
for adapter in self._adapters
|
||
if str(getattr(adapter, "adapter_name", "") or "").strip()
|
||
}
|
||
return self.default_provider_statuses(
|
||
mcp_available="mcp" in adapter_names,
|
||
live_search_available="live_search" in adapter_names,
|
||
)
|
||
|
||
@classmethod
|
||
def default_provider_statuses(
|
||
cls,
|
||
*,
|
||
mcp_available: bool | None = None,
|
||
live_search_available: bool | None = None,
|
||
) -> list[dict[str, Any]]:
|
||
mcp_enabled = _env_enabled("AI_EXTERNAL_QUERY_MCP_ADAPTER_ENABLED", True)
|
||
live_search_enabled = _env_enabled("AI_EXTERNAL_QUERY_LIVE_SEARCH_FALLBACK_ENABLED", True)
|
||
if mcp_available is not None:
|
||
mcp_enabled = bool(mcp_available)
|
||
if live_search_available is not None:
|
||
live_search_enabled = bool(live_search_available)
|
||
return [
|
||
cls._provider_status("mcp", available=mcp_enabled),
|
||
cls._provider_status("live_search", available=live_search_enabled),
|
||
]
|
||
|
||
@staticmethod
|
||
def _provider_status(provider_id: str, *, available: bool) -> dict[str, Any]:
|
||
spec = dict(EXTERNAL_QUERY_PROVIDER_SPECS.get(provider_id) or {})
|
||
spec.setdefault("id", provider_id)
|
||
spec.setdefault("label", provider_id)
|
||
spec["available"] = bool(available)
|
||
spec["status"] = "available" if available else "unavailable"
|
||
return spec
|
||
|
||
@classmethod
|
||
def status_snapshot(cls, *, binding: dict[str, Any] | None = None) -> dict[str, Any]:
|
||
providers = cls.default_provider_statuses()
|
||
provider_by_id = {str(item.get("id") or ""): item for item in providers}
|
||
raw_binding = dict(binding or {}) if isinstance(binding, dict) else {}
|
||
primary_provider_id = str(raw_binding.get("primary_provider_id") or "").strip()
|
||
fallback_provider_ids = [
|
||
str(item or "").strip()
|
||
for item in list(raw_binding.get("fallback_provider_ids") or [])
|
||
if str(item or "").strip()
|
||
]
|
||
verifier_available = _env_enabled("AI_EXTERNAL_QUERY_VERIFIER_ENABLED", True)
|
||
source_filtering_fail_closed = _env_enabled("AI_EXTERNAL_QUERY_SOURCE_FILTERING_FAIL_CLOSED", True)
|
||
selected_ids = [primary_provider_id, *fallback_provider_ids]
|
||
selected_available = [
|
||
provider_by_id[item]
|
||
for item in selected_ids
|
||
if item in provider_by_id and bool(provider_by_id[item].get("available"))
|
||
]
|
||
return {
|
||
"providers": providers,
|
||
"provider_ids": [str(item.get("id") or "") for item in providers],
|
||
"available_provider_count": len([item for item in providers if bool(item.get("available"))]),
|
||
"selected_available_provider_count": len(selected_available),
|
||
"primary_provider_id": primary_provider_id or None,
|
||
"fallback_provider_ids": fallback_provider_ids,
|
||
"mcp_adapter_available": bool(provider_by_id.get("mcp", {}).get("available")),
|
||
"live_search_fallback_available": bool(provider_by_id.get("live_search", {}).get("available")),
|
||
"verifier": {
|
||
"required": True,
|
||
"available": verifier_available,
|
||
"status": "available" if verifier_available else "unavailable",
|
||
},
|
||
"source_filtering": {
|
||
"required": True,
|
||
"fail_closed": source_filtering_fail_closed,
|
||
"status": "fail_closed" if source_filtering_fail_closed else "open",
|
||
},
|
||
}
|