319 lines
13 KiB
Python
319 lines
13 KiB
Python
"""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
|
|
import hashlib
|
|
import json
|
|
from typing import Any, Protocol
|
|
from uuid import uuid4
|
|
|
|
from .search_profile_resolver import DomainQueryProfileRegistry, SearchProfile, SearchProfileResolver
|
|
from .trace_tool_credential_governance import (
|
|
GovernanceActor,
|
|
GovernanceError,
|
|
GovernanceStore,
|
|
TraceContext,
|
|
)
|
|
|
|
|
|
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, *, governance_store: GovernanceStore | None = None):
|
|
self.providers = providers or []
|
|
self.governance_store = governance_store
|
|
self.profile_resolver = SearchProfileResolver()
|
|
self.planner = QueryPlanner()
|
|
self.guard = SearchPolicyGuard()
|
|
|
|
def search(self, request: SearchRequest, *, actor: GovernanceActor | dict[str, Any] | None = None, trace: TraceContext | None = None) -> SearchResponse:
|
|
trusted_actor: GovernanceActor | None = None
|
|
if self.governance_store is not None:
|
|
if actor is None or trace is None:
|
|
raise GovernanceError("TOOL_MCP_AUTHZ_DENY")
|
|
trusted_actor = GovernanceActor.from_value(actor)
|
|
if request.requester != trusted_actor.identity:
|
|
raise GovernanceError("TOOL_MCP_AUTHZ_DENY", field="requester")
|
|
if request.correlation_id is not None and request.correlation_id not in {
|
|
trace.trace_id, trace.correlation_id
|
|
}:
|
|
raise GovernanceError("TOOL_SCOPE_PROPAGATION", field="correlation_id")
|
|
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:
|
|
if request.authorized_sources and provider.name not in request.authorized_sources:
|
|
attempts.append(SearchProviderAttempt(provider=provider.name, status="denied", reason="SOURCE_AUTHZ_DENY"))
|
|
continue
|
|
if request.user_scope != "public" and provider.name not in request.authorized_sources:
|
|
attempts.append(SearchProviderAttempt(provider=provider.name, status="denied", reason="SOURCE_AUTHZ_REQUIRED"))
|
|
continue
|
|
try:
|
|
if self.governance_store is not None:
|
|
assert trusted_actor is not None and trace is not None
|
|
input_digest = hashlib.sha256(
|
|
json.dumps(
|
|
{"query": request.query, "intent": request.intent, "provider": provider.name},
|
|
sort_keys=True, separators=(",", ":"), ensure_ascii=False,
|
|
).encode("utf-8")
|
|
).hexdigest()
|
|
invocation_key = hashlib.sha256(
|
|
f"{trace.trace_id}:{trace.span_id}:{provider.name}:{input_digest}".encode("utf-8")
|
|
).hexdigest()
|
|
authorization = self.governance_store.authorize_tool(
|
|
actor=trusted_actor,
|
|
tool_name=f"search.{provider.name}",
|
|
trace=trace,
|
|
requested_scope=trace.scope,
|
|
input_digest=input_digest,
|
|
invocation_key=invocation_key,
|
|
)
|
|
if authorization.get("idempotent"):
|
|
raise GovernanceError("TOOL_INVOCATION_REPLAY")
|
|
provider_evidence, attempt = provider.search(request, expanded_queries)
|
|
if self.governance_store is not None:
|
|
result_digest = hashlib.sha256(
|
|
json.dumps(
|
|
{
|
|
"evidence": [item.to_dict() for item in provider_evidence],
|
|
"attempt": attempt.to_dict(),
|
|
},
|
|
sort_keys=True, separators=(",", ":"), ensure_ascii=False,
|
|
).encode("utf-8")
|
|
).hexdigest()
|
|
self.governance_store.record_tool_receipt(
|
|
invocation_key=invocation_key,
|
|
result_digest=result_digest,
|
|
actor=trusted_actor,
|
|
trace=trace,
|
|
)
|
|
evidence.extend(provider_evidence)
|
|
attempts.append(attempt)
|
|
if provider_evidence:
|
|
break
|
|
except GovernanceError:
|
|
raise
|
|
except Exception:
|
|
# Provider exceptions may contain URLs, tokens, or local paths.
|
|
attempts.append(SearchProviderAttempt(provider=getattr(provider, "name", "unknown"), status="failed", reason="PROVIDER_FAILURE"))
|
|
|
|
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",
|
|
]
|