PMFlow/backend/app/services/external_query_service/registry.py

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 与来源过滤。",
    },
    "dashscope_responses": {
        "id": "dashscope_responses",
        "label": "DashScope Responses 内置搜索",
        "kind": "model_builtin_web_search",
        "summary": "通过 qwen3.6-plus Responses API 内置 web_search / web_extractor 获取公开来源。",
    },
}


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,
            dashscope_responses_available="dashscope_responses" in adapter_names,
        )

    @classmethod
    def default_provider_statuses(
        cls,
        *,
        mcp_available: bool | None = None,
        live_search_available: bool | None = None,
        dashscope_responses_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)
        dashscope_responses_enabled = _env_enabled("AI_EXTERNAL_QUERY_DASHSCOPE_RESPONSES_ENABLED", True) and bool(
            str(os.getenv("DASHSCOPE_API_KEY") or "").strip()
        )
        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)
        if dashscope_responses_available is not None:
            dashscope_responses_enabled = bool(dashscope_responses_available)
        return [
            cls._provider_status("dashscope_responses", available=dashscope_responses_enabled),
            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")),
            "dashscope_responses_available": bool(provider_by_id.get("dashscope_responses", {}).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",
            },
        }