355 lines
15 KiB
Python
355 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import inspect
|
|
import os
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from .gateway import ExternalQueryGateway
|
|
from .registry import ExternalProviderRegistry
|
|
from .schemas import ExternalAdapterResult, ExternalEvidenceItem, ExternalQueryExecution, ExternalQueryPlan
|
|
|
|
|
|
@dataclass
|
|
class _ProviderHealth:
|
|
consecutive_failures: int = 0
|
|
opened_until: float = 0.0
|
|
|
|
|
|
class ExternalProviderHealthRegistry:
|
|
"""In-process circuit breaker for external providers.
|
|
|
|
This is deliberately runtime-level state, not business state. Its only job is
|
|
to stop a repeatedly failing provider from adding latency to every user turn.
|
|
"""
|
|
|
|
_state: dict[str, _ProviderHealth] = {}
|
|
|
|
@classmethod
|
|
def is_open(cls, provider_id: str) -> bool:
|
|
state = cls._state.get(str(provider_id or ""))
|
|
return bool(state and state.opened_until > time.monotonic())
|
|
|
|
@classmethod
|
|
def record_success(cls, provider_id: str) -> None:
|
|
provider = str(provider_id or "").strip()
|
|
if not provider:
|
|
return
|
|
cls._state[provider] = _ProviderHealth()
|
|
|
|
@classmethod
|
|
def record_failure(cls, provider_id: str) -> None:
|
|
provider = str(provider_id or "").strip()
|
|
if not provider:
|
|
return
|
|
state = cls._state.setdefault(provider, _ProviderHealth())
|
|
state.consecutive_failures += 1
|
|
threshold = _env_int("AI_EXTERNAL_PROVIDER_CIRCUIT_FAILURE_THRESHOLD", 2)
|
|
cooldown = _env_float("AI_EXTERNAL_PROVIDER_CIRCUIT_COOLDOWN_SECONDS", 180.0)
|
|
if state.consecutive_failures >= max(1, threshold):
|
|
state.opened_until = time.monotonic() + max(5.0, cooldown)
|
|
|
|
|
|
class ExternalResearchRuntime:
|
|
"""Agentic external research executor.
|
|
|
|
The old gateway executes adapters one by one. This runtime builds one
|
|
external plan, fans it out to available providers in parallel, hydrates and
|
|
verifies evidence, and returns as soon as enough verified coverage exists.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
planner: Any,
|
|
registry: ExternalProviderRegistry,
|
|
adapter_timeout_seconds: float | None = None,
|
|
max_parallel_adapters: int | None = None,
|
|
min_verified_sources: int | None = None,
|
|
) -> None:
|
|
self.planner = planner
|
|
self.registry = registry
|
|
self.adapter_timeout_seconds = _normalize_timeout(
|
|
adapter_timeout_seconds,
|
|
default=_env_float("AI_EXTERNAL_RESEARCH_ADAPTER_TIMEOUT_SECONDS", 14.0),
|
|
)
|
|
self.max_parallel_adapters = max(1, int(max_parallel_adapters or _env_int("AI_EXTERNAL_RESEARCH_MAX_PARALLEL_ADAPTERS", 3)))
|
|
self.min_verified_sources = max(1, int(min_verified_sources or _env_int("AI_EXTERNAL_RESEARCH_MIN_VERIFIED_SOURCES", 2)))
|
|
|
|
async def execute(
|
|
self,
|
|
*,
|
|
question: str,
|
|
query_semantics: dict[str, Any] | None = None,
|
|
fit_brief: dict[str, Any] | None = None,
|
|
) -> ExternalQueryExecution:
|
|
plan, planner_meta = await self._build_plan(
|
|
question=question,
|
|
query_semantics=query_semantics,
|
|
fit_brief=fit_brief,
|
|
)
|
|
ranked_adapters = self._rank_adapters(plan=plan)
|
|
adapters, circuit_attempts = self._apply_circuit_breaker(ranked_adapters)
|
|
if not adapters and ranked_adapters:
|
|
adapters = ranked_adapters[:1]
|
|
|
|
attempts: list[dict[str, Any]] = list(circuit_attempts)
|
|
provider_results: list[tuple[ExternalAdapterResult, dict[str, Any]]] = []
|
|
pending: dict[asyncio.Task[tuple[ExternalAdapterResult, dict[str, Any], dict[str, Any]]], str] = {}
|
|
for adapter in adapters[: self.max_parallel_adapters]:
|
|
adapter_name = self._adapter_name(adapter)
|
|
pending[asyncio.create_task(self._run_adapter(adapter=adapter, plan=plan))] = adapter_name
|
|
|
|
selected_evidence: list[ExternalEvidenceItem] = []
|
|
selected_verifier = ExternalQueryGateway._verify_evidence_bundle(plan=plan, evidence=[])
|
|
started = time.perf_counter()
|
|
while pending:
|
|
done, remaining = await asyncio.wait(set(pending), return_when=asyncio.FIRST_COMPLETED)
|
|
for task in done:
|
|
adapter_name = pending.pop(task, "")
|
|
try:
|
|
result, verifier, attempt = task.result()
|
|
except Exception as exc:
|
|
attempt = {
|
|
"adapter": adapter_name or "unknown",
|
|
"status": "error",
|
|
"evidence_count": 0,
|
|
"verifier_ok": False,
|
|
"verifier_reason": "adapter_task_error",
|
|
"error_type": type(exc).__name__,
|
|
}
|
|
attempts.append(attempt)
|
|
ExternalProviderHealthRegistry.record_failure(adapter_name)
|
|
continue
|
|
attempts.append(attempt)
|
|
provider_results.append((result, verifier))
|
|
selected_evidence = self._rank_evidence(plan=plan, provider_results=provider_results)
|
|
selected_verifier = ExternalQueryGateway._verify_evidence_bundle(plan=plan, evidence=selected_evidence)
|
|
if self._has_enough_verified_coverage(plan=plan, evidence=selected_evidence, verifier=selected_verifier):
|
|
for pending_task in remaining:
|
|
pending_task.cancel()
|
|
pending.clear()
|
|
break
|
|
|
|
if not selected_evidence:
|
|
selected_evidence = self._rank_evidence(plan=plan, provider_results=provider_results)
|
|
selected_verifier = ExternalQueryGateway._verify_evidence_bundle(plan=plan, evidence=selected_evidence)
|
|
|
|
best_adapter = self._best_adapter_name(provider_results=provider_results)
|
|
meta = {
|
|
"fallback_used": False,
|
|
"adapter_attempts": attempts,
|
|
"selected_status": "success" if selected_evidence else "empty",
|
|
"provider_results": [
|
|
{
|
|
"adapter": result.adapter_name,
|
|
"status": result.status,
|
|
"evidence_count": len(result.evidence or []),
|
|
"verifier_ok": bool(verifier.get("ok")),
|
|
"verifier_reason": verifier.get("reason_code"),
|
|
}
|
|
for result, verifier in provider_results
|
|
],
|
|
"research_runtime": "parallel_external_research.v1",
|
|
"research_latency_ms": int((time.perf_counter() - started) * 1000),
|
|
"query_count": len(plan.search_queries or []),
|
|
"verified_source_count": len(selected_evidence or []),
|
|
}
|
|
return ExternalQueryExecution(
|
|
plan=plan,
|
|
adapter_name=best_adapter or "external_research",
|
|
evidence=selected_evidence,
|
|
verifier=selected_verifier,
|
|
planner_meta=planner_meta,
|
|
meta=meta,
|
|
)
|
|
|
|
async def _build_plan(
|
|
self,
|
|
*,
|
|
question: str,
|
|
query_semantics: dict[str, Any] | None,
|
|
fit_brief: dict[str, Any] | None,
|
|
) -> tuple[ExternalQueryPlan, dict[str, Any]]:
|
|
build_plan = getattr(self.planner, "build_plan")
|
|
try:
|
|
parameters = inspect.signature(build_plan).parameters
|
|
except Exception:
|
|
parameters = {}
|
|
kwargs = {"question": question, "query_semantics": query_semantics}
|
|
if "fit_brief" in parameters:
|
|
kwargs["fit_brief"] = fit_brief
|
|
planned = await build_plan(**kwargs)
|
|
planner_meta: dict[str, Any] = {}
|
|
if isinstance(planned, tuple) and len(planned) == 2:
|
|
plan, planner_meta = planned
|
|
else:
|
|
plan = planned
|
|
if not isinstance(plan, ExternalQueryPlan):
|
|
plan = ExternalQueryPlan.model_validate(plan)
|
|
return plan, dict(planner_meta or {}) if isinstance(planner_meta, dict) else {}
|
|
|
|
def _rank_adapters(self, *, plan: ExternalQueryPlan) -> list[Any]:
|
|
gateway = ExternalQueryGateway(planner=self.planner, registry=self.registry)
|
|
return gateway._rank_adapters(plan=plan)
|
|
|
|
@classmethod
|
|
def _apply_circuit_breaker(cls, adapters: list[Any]) -> tuple[list[Any], list[dict[str, Any]]]:
|
|
active: list[Any] = []
|
|
skipped: list[dict[str, Any]] = []
|
|
for adapter in adapters:
|
|
adapter_name = cls._adapter_name(adapter)
|
|
if ExternalProviderHealthRegistry.is_open(adapter_name):
|
|
skipped.append(
|
|
{
|
|
"adapter": adapter_name,
|
|
"status": "circuit_open",
|
|
"evidence_count": 0,
|
|
"verifier_ok": False,
|
|
"verifier_reason": "provider_recent_failures",
|
|
}
|
|
)
|
|
continue
|
|
active.append(adapter)
|
|
return active, skipped
|
|
|
|
async def _run_adapter(
|
|
self,
|
|
*,
|
|
adapter: Any,
|
|
plan: ExternalQueryPlan,
|
|
) -> tuple[ExternalAdapterResult, dict[str, Any], dict[str, Any]]:
|
|
adapter_name = self._adapter_name(adapter)
|
|
started = time.perf_counter()
|
|
try:
|
|
result = await asyncio.wait_for(adapter.execute(plan), timeout=self.adapter_timeout_seconds)
|
|
except asyncio.TimeoutError:
|
|
ExternalProviderHealthRegistry.record_failure(adapter_name)
|
|
result = ExternalAdapterResult(adapter_name=adapter_name, status="timeout", evidence=[], meta={})
|
|
verifier = ExternalQueryGateway._verify_evidence_bundle(plan=plan, evidence=[])
|
|
return result, verifier, {
|
|
"adapter": adapter_name,
|
|
"status": "timeout",
|
|
"evidence_count": 0,
|
|
"verifier_ok": False,
|
|
"verifier_reason": "adapter_timeout",
|
|
"latency_ms": int((time.perf_counter() - started) * 1000),
|
|
"timeout_seconds": self.adapter_timeout_seconds,
|
|
}
|
|
except Exception as exc:
|
|
ExternalProviderHealthRegistry.record_failure(adapter_name)
|
|
result = ExternalAdapterResult(adapter_name=adapter_name, status="error", evidence=[], meta={"error_type": type(exc).__name__})
|
|
verifier = ExternalQueryGateway._verify_evidence_bundle(plan=plan, evidence=[])
|
|
return result, verifier, {
|
|
"adapter": adapter_name,
|
|
"status": "error",
|
|
"evidence_count": 0,
|
|
"verifier_ok": False,
|
|
"verifier_reason": "adapter_error",
|
|
"error_type": type(exc).__name__,
|
|
"latency_ms": int((time.perf_counter() - started) * 1000),
|
|
}
|
|
if not isinstance(result, ExternalAdapterResult):
|
|
result = ExternalAdapterResult.model_validate(result)
|
|
verifier = ExternalQueryGateway._verify_evidence_bundle(plan=plan, evidence=list(result.evidence or []))
|
|
ok = bool(result.evidence) and bool(verifier.get("ok"))
|
|
if ok:
|
|
ExternalProviderHealthRegistry.record_success(adapter_name)
|
|
elif str(result.status or "").strip().lower() in {"error", "timeout"}:
|
|
ExternalProviderHealthRegistry.record_failure(adapter_name)
|
|
return result, verifier, {
|
|
"adapter": adapter_name,
|
|
"status": str(result.status or "").strip() or "empty",
|
|
"evidence_count": len(result.evidence or []),
|
|
"verifier_ok": bool(verifier.get("ok")),
|
|
"verifier_reason": verifier.get("reason_code"),
|
|
"latency_ms": int((time.perf_counter() - started) * 1000),
|
|
}
|
|
|
|
@classmethod
|
|
def _rank_evidence(
|
|
cls,
|
|
*,
|
|
plan: ExternalQueryPlan,
|
|
provider_results: list[tuple[ExternalAdapterResult, dict[str, Any]]],
|
|
) -> list[ExternalEvidenceItem]:
|
|
rows: list[tuple[float, ExternalEvidenceItem]] = []
|
|
seen: set[str] = set()
|
|
for result, verifier in provider_results:
|
|
if not result.evidence:
|
|
continue
|
|
verifier_ok = bool(verifier.get("ok"))
|
|
for item in list(result.evidence or []):
|
|
key = cls._evidence_key(item)
|
|
if key and key in seen:
|
|
continue
|
|
if key:
|
|
seen.add(key)
|
|
score = max(float(item.authority_score or 0.0), float(item.freshness_score or 0.0))
|
|
if verifier_ok:
|
|
score += 0.25
|
|
if str(item.target_url or "").startswith("https://"):
|
|
score += 0.03
|
|
rows.append((round(score, 4), item))
|
|
rows.sort(key=lambda item: item[0], reverse=True)
|
|
return [item for _, item in rows[: max(1, int(plan.result_limit or 3))]]
|
|
|
|
def _has_enough_verified_coverage(
|
|
self,
|
|
*,
|
|
plan: ExternalQueryPlan,
|
|
evidence: list[ExternalEvidenceItem],
|
|
verifier: dict[str, Any],
|
|
) -> bool:
|
|
if not evidence or not bool(verifier.get("ok")):
|
|
return False
|
|
required = 1 if not bool(plan.requires_multi_source) else 2
|
|
required = max(required, min(3, int(plan.result_limit or required), self.min_verified_sources))
|
|
return len(evidence) >= max(1, required)
|
|
|
|
@staticmethod
|
|
def _best_adapter_name(provider_results: list[tuple[ExternalAdapterResult, dict[str, Any]]]) -> str:
|
|
for result, verifier in provider_results:
|
|
if result.evidence and bool(verifier.get("ok")):
|
|
return str(result.adapter_name or "").strip() or "external_research"
|
|
for result, _verifier in provider_results:
|
|
if result.evidence:
|
|
return str(result.adapter_name or "").strip() or "external_research"
|
|
return "external_research"
|
|
|
|
@staticmethod
|
|
def _adapter_name(adapter: Any) -> str:
|
|
return str(getattr(adapter, "adapter_name", "") or adapter.__class__.__name__ or "adapter").strip()
|
|
|
|
@staticmethod
|
|
def _evidence_key(item: ExternalEvidenceItem) -> str:
|
|
url = str(item.target_url or "").strip().lower()
|
|
if url:
|
|
return f"url:{url}"
|
|
title = str(item.title or item.source_label or item.snippet or "").strip().lower()
|
|
return f"title:{title[:160]}" if title else ""
|
|
|
|
|
|
def _normalize_timeout(value: float | None, *, default: float) -> float:
|
|
try:
|
|
raw = float(value if value is not None else default)
|
|
except Exception:
|
|
raw = float(default)
|
|
return max(3.0, min(raw, 30.0))
|
|
|
|
|
|
def _env_float(name: str, default: float) -> float:
|
|
try:
|
|
return float(str(os.getenv(name) or "").strip() or default)
|
|
except Exception:
|
|
return default
|
|
|
|
|
|
def _env_int(name: str, default: int) -> int:
|
|
try:
|
|
return int(str(os.getenv(name) or "").strip() or default)
|
|
except Exception:
|
|
return default
|