aios/active-runtime/app/search_gateway.py

"""Governed general-purpose Search Gateway for AIOS.

The gateway coordinates query planning, provider routing, fallback, evidence
normalization and compliance checks. It is provider-agnostic: runtime workers plug
in public or explicitly authorized search providers.

Search is a shared capability. Groups and users may customize provider order,
trusted domains, query expansion hints and domain profiles, but customization may
not override hard authorization, privacy, safety or audit rules.
"""

from __future__ import annotations

from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from typing import Any, Protocol
from uuid import uuid4

from .search_profile_resolver import DomainQueryProfileRegistry, SearchProfile, SearchProfileResolver


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


@dataclass(frozen=True)
class SearchRequest:
    query: str
    requester: str
    intent: str = "general"
    region: str | None = None
    time_range: str | None = None
    data_type: str | None = None
    user_scope: str = "public"
    authorized_sources: list[str] = field(default_factory=list)
    domain_profiles: list[str] = field(default_factory=list)
    group_profile: dict[str, Any] | None = None
    user_profile: dict[str, Any] | None = None
    session_profile: dict[str, Any] | None = None
    correlation_id: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True)
class SearchEvidence:
    title: str
    url: str | None
    source: str
    snippet: str
    evidence_type: str = "web"
    confidence: str = "medium"
    retrieved_at: str = field(default_factory=utc_now)

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True)
class SearchProviderAttempt:
    provider: str
    status: str
    result_count: int = 0
    reason: str = ""

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True)
class SearchResponse:
    request_id: str
    status: str
    answer: str
    evidence: list[SearchEvidence]
    attempts: list[SearchProviderAttempt]
    confidence: str
    expanded_queries: list[str] = field(default_factory=list)
    resolved_profile: dict[str, Any] | None = None
    need_clarification: list[str] = field(default_factory=list)
    blocked_reasons: list[str] = field(default_factory=list)
    created_at: str = field(default_factory=utc_now)

    def to_dict(self) -> dict[str, Any]:
        return {
            "request_id": self.request_id,
            "status": self.status,
            "answer": self.answer,
            "evidence": [e.to_dict() for e in self.evidence],
            "attempts": [a.to_dict() for a in self.attempts],
            "confidence": self.confidence,
            "expanded_queries": self.expanded_queries,
            "resolved_profile": self.resolved_profile,
            "need_clarification": self.need_clarification,
            "blocked_reasons": self.blocked_reasons,
            "created_at": self.created_at,
        }


class SearchProvider(Protocol):
    name: str

    def search(self, request: SearchRequest, queries: list[str]) -> tuple[list[SearchEvidence], SearchProviderAttempt]:
        ...


class QueryPlanner:
    def __init__(self, domain_registry: DomainQueryProfileRegistry | None = None):
        self.domain_registry = domain_registry or DomainQueryProfileRegistry()

    def expand(self, request: SearchRequest, profile: SearchProfile) -> list[str]:
        base = request.query.strip()
        queries = [base]
        if request.region:
            queries.append(f"{request.region} {base}")
        if request.time_range:
            queries.append(f"{base} {request.time_range}")
        domain_profiles = list(dict.fromkeys([request.intent] + request.domain_profiles + profile.domain_profiles))
        hints = list(profile.query_expansion_hints) + self.domain_registry.hints_for(domain_profiles)
        for hint in hints:
            queries.append(f"{base} {hint}")
        for domain in profile.trusted_domains:
            queries.append(f"site:{domain} {base}")
        return list(dict.fromkeys([q for q in queries if q]))


class SearchPolicyGuard:
    """Blocks non-public or unauthorized collection patterns."""

    FORBIDDEN_TERMS = (
        "绕过登录",
        "绕过验证码",
        "验证码绕过",
        "破解登录",
        "个人成绩",
        "学生隐私",
        "受保护平台批量爬取",
    )

    def evaluate(self, request: SearchRequest, profile: SearchProfile) -> list[str]:
        q = request.query.lower()
        reasons: list[str] = []
        for term in self.FORBIDDEN_TERMS:
            if term.lower() in q:
                reasons.append(f"blocked_term:{term}")
        if request.user_scope not in {"public", "authorized_user_file", "authorized_link", "internal_archive"}:
            reasons.append(f"unsupported_scope:{request.user_scope}")
        if request.data_type in {"personal_grade", "student_private_data", "private_person_data"} and request.user_scope != "authorized_user_file":
            reasons.append("private_data_requires_user_authorized_file")
        if not profile.audit_required:
            reasons.append("audit_cannot_be_disabled")
        if not profile.fallback_required:
            reasons.append("fallback_cannot_be_disabled")
        return reasons


class SearchGateway:
    def __init__(self, providers: list[SearchProvider] | None = None):
        self.providers = providers or []
        self.profile_resolver = SearchProfileResolver()
        self.planner = QueryPlanner()
        self.guard = SearchPolicyGuard()

    def search(self, request: SearchRequest) -> SearchResponse:
        request_id = request.correlation_id or f"search_{uuid4().hex[:12]}"
        profile = self.profile_resolver.resolve(request.group_profile, request.user_profile, request.session_profile)
        expanded_queries = self.planner.expand(request, profile)
        blocked = self.guard.evaluate(request, profile)
        if blocked:
            return SearchResponse(
                request_id=request_id,
                status="blocked",
                answer="该搜索请求涉及未授权访问、验证码/登录绕过、受保护平台批量抓取或个人隐私数据，已阻断。可改为搜索公开网页、官方公告、公开文档，或处理用户授权上传的文件/链接。",
                evidence=[],
                attempts=[],
                confidence="high",
                expanded_queries=expanded_queries,
                resolved_profile=profile.to_dict(),
                blocked_reasons=blocked,
            )

        evidence: list[SearchEvidence] = []
        attempts: list[SearchProviderAttempt] = []
        providers = self._ordered_providers(profile)
        for provider in providers:
            try:
                provider_evidence, attempt = provider.search(request, expanded_queries)
                evidence.extend(provider_evidence)
                attempts.append(attempt)
                if provider_evidence:
                    break
            except Exception as exc:
                attempts.append(SearchProviderAttempt(provider=getattr(provider, "name", "unknown"), status="failed", reason=str(exc)))

        if evidence:
            return SearchResponse(
                request_id=request_id,
                status="found",
                answer="已找到公开或授权来源结果。",
                evidence=evidence,
                attempts=attempts,
                confidence=self._confidence(evidence),
                expanded_queries=expanded_queries,
                resolved_profile=profile.to_dict(),
            )

        return SearchResponse(
            request_id=request_id,
            status="not_found",
            answer="未找到足够可靠的公开或授权结果；已记录 provider 尝试和失败原因。",
            evidence=[],
            attempts=attempts,
            confidence="low",
            expanded_queries=expanded_queries,
            resolved_profile=profile.to_dict(),
            need_clarification=self._clarification(request),
        )

    def _ordered_providers(self, profile: SearchProfile) -> list[SearchProvider]:
        if not self.providers:
            return []
        by_name = {provider.name: provider for provider in self.providers}
        ordered = [by_name[name] for name in profile.provider_order if name in by_name]
        remaining = [provider for provider in self.providers if provider.name not in {p.name for p in ordered}]
        return ordered + remaining

    def _confidence(self, evidence: list[SearchEvidence]) -> str:
        if any(e.confidence == "high" for e in evidence):
            return "high"
        if evidence:
            return "medium"
        return "low"

    def _clarification(self, request: SearchRequest) -> list[str]:
        if request.intent == "education_exam_data":
            return ["请补充学校名称、联考名称、年级、科目，以及需要试卷/答案还是成绩统计。"]
        if request.intent in {"company_research", "policy_regulation", "software_library", "product_research"}:
            return ["请补充主体名称、地区/版本、时间范围、需要的数据类型或可信来源范围。"]
        return ["请补充更具体的主体、平台、地区、时间、数据类型或授权文件/链接。"]


__all__ = [
    "SearchRequest",
    "SearchEvidence",
    "SearchProviderAttempt",
    "SearchResponse",
    "SearchProvider",
    "QueryPlanner",
    "SearchPolicyGuard",
    "SearchGateway",
]