agent_search_gateway.scheduler.fetch

  1"""Sequential URL-fetch provider scheduler with typed outcomes."""
  2
  3import asyncio
  4import logging
  5import time
  6from collections.abc import Callable, Sequence
  7
  8from ..concurrency import CapacityLease, ProviderQuotaManager
  9from ..errors import ErrorCode, ExecutionFailure
 10from ..llm.stages import LLMStages, cheap_check
 11from ..models import FetchOutcome
 12from ..observability import elapsed_ms, log_event, target_url_for_log
 13from ..providers.contracts import URLFetchCandidate, URLFetchProvider
 14from ..url_normalization import NormalizedURL
 15
 16
 17class FetchScheduler:
 18    def __init__(
 19        self,
 20        providers: Sequence[URLFetchProvider],
 21        quotas: ProviderQuotaManager,
 22        stages: LLMStages,
 23        *,
 24        logger: logging.Logger | None = None,
 25        monotonic: Callable[[], float] = time.monotonic,
 26    ) -> None:
 27        self._providers = tuple(providers)
 28        self._quotas = quotas
 29        self._stages = stages
 30        self._logger = logger or logging.getLogger(__name__)
 31        self._monotonic = monotonic
 32
 33    @property
 34    def provider_names(self) -> tuple[str, ...]:
 35        return tuple(provider.name for provider in self._providers)
 36
 37    async def fetch_until_accepted(self, url: NormalizedURL) -> FetchOutcome:
 38        semantic_failure_seen = False
 39        failures: list[ExecutionFailure] = []
 40        remaining = list(self._providers)
 41        while remaining:
 42            provider, lease = await self._select_available_provider(remaining)
 43            remaining.remove(provider)
 44            log_event(
 45                self._logger,
 46                logging.DEBUG,
 47                "provider_selected",
 48                provider=provider.name,
 49                url=target_url_for_log(str(url)),
 50                candidate_count=len(remaining) + 1,
 51            )
 52            try:
 53                async with lease:
 54                    outcome = await self._attempt(provider, url)
 55            except asyncio.CancelledError:
 56                raise
 57            if outcome.kind == "accepted":
 58                return FetchOutcome("accepted", outcome.candidate, tuple(failures))
 59            log_event(
 60                self._logger,
 61                logging.DEBUG,
 62                "provider_fallback",
 63                provider=provider.name,
 64                url=target_url_for_log(str(url)),
 65                outcome=outcome.kind,
 66                remaining=len(remaining),
 67            )
 68            if outcome.kind == "semantic_failure":
 69                semantic_failure_seen = True
 70            else:
 71                failures.extend(outcome.failures)
 72
 73        if semantic_failure_seen:
 74            return FetchOutcome("semantic_failure", failures=tuple(failures))
 75        return FetchOutcome("execution_failure", failures=tuple(failures))
 76
 77    async def _select_available_provider(
 78        self,
 79        remaining: list[URLFetchProvider],
 80    ) -> tuple[URLFetchProvider, CapacityLease]:
 81        while True:
 82            for provider in remaining:
 83                lease = await self._quotas.get_web(provider.name).try_lease()
 84                if lease is not None:
 85                    return provider, lease
 86            log_event(
 87                self._logger,
 88                logging.DEBUG,
 89                "scheduler_waiting",
 90                candidate_count=len(remaining),
 91            )
 92            await self._quotas.wait_until_any_web_available(
 93                tuple(provider.name for provider in remaining)
 94            )
 95
 96    async def _attempt(
 97        self,
 98        provider: URLFetchProvider,
 99        url: NormalizedURL,
100    ) -> FetchOutcome:
101        started = self._monotonic()
102        log_event(
103            self._logger,
104            logging.DEBUG,
105            "provider_started",
106            provider=provider.name,
107            stage="fetch",
108            url=target_url_for_log(str(url)),
109        )
110        try:
111            candidate = await provider.fetch(url)
112            self._validate_candidate(candidate)
113        except asyncio.CancelledError:
114            raise
115        except ExecutionFailure as exc:
116            self._log_provider_failed(provider.name, url, started, exc)
117            return FetchOutcome("execution_failure", failures=(exc,))
118        except Exception as exc:
119            self._log_provider_failed(provider.name, url, started, exc)
120            failure = ExecutionFailure(
121                ErrorCode.ALL_PROVIDERS_FAILED,
122                f"Fetch provider {provider.name} returned invalid data",
123            )
124            return FetchOutcome("execution_failure", failures=(failure,))
125
126        log_event(
127            self._logger,
128            logging.DEBUG,
129            "candidate_accepted",
130            provider=provider.name,
131            url=target_url_for_log(str(url)),
132            raw_chars=len(candidate.raw_content),
133            content_chars=len(candidate.content),
134        )
135        validation_candidate = candidate.content or candidate.raw_content
136        if not cheap_check(validation_candidate):
137            self._log_body_rejected(provider.name, url, "cheap_check")
138            self._log_provider_completed(provider.name, url, started, "semantic_failure")
139            return FetchOutcome("semantic_failure")
140
141        try:
142            decision = await self._stages.judge(validation_candidate)
143        except asyncio.CancelledError:
144            raise
145        except ExecutionFailure as exc:
146            failure = ExecutionFailure(
147                exc.code,
148                exc.message,
149                reason=exc.reason or "judge_execution_failed",
150            )
151            self._log_provider_failed(
152                provider.name,
153                url,
154                started,
155                failure,
156                failed_stage="judge",
157                dependency=self._stages.judge_provider,
158            )
159            raise failure from exc
160        except Exception as exc:
161            failure = ExecutionFailure(
162                ErrorCode.LLM_STAGE_FAILED,
163                "Judge LLM stage failed",
164                reason="judge_execution_failed",
165            )
166            self._log_provider_failed(
167                provider.name,
168                url,
169                started,
170                failure,
171                failed_stage="judge",
172                dependency=self._stages.judge_provider,
173            )
174            raise failure from exc
175
176        if not decision.ok:
177            self._log_body_rejected(provider.name, url, "judge_rejected")
178            self._log_provider_completed(provider.name, url, started, "semantic_failure")
179            return FetchOutcome("semantic_failure")
180        log_event(
181            self._logger,
182            logging.DEBUG,
183            "body_accepted",
184            provider=provider.name,
185            url=target_url_for_log(str(url)),
186            raw_chars=len(candidate.raw_content),
187            content_chars=len(candidate.content),
188        )
189        self._log_provider_completed(provider.name, url, started, "accepted")
190        return FetchOutcome("accepted", candidate)
191
192    def _log_body_rejected(
193        self,
194        provider: str,
195        url: NormalizedURL,
196        reason: str,
197    ) -> None:
198        log_event(
199            self._logger,
200            logging.DEBUG,
201            "body_rejected",
202            provider=provider,
203            url=target_url_for_log(str(url)),
204            reason=reason,
205        )
206
207    def _log_provider_completed(
208        self,
209        provider: str,
210        url: NormalizedURL,
211        started: float,
212        outcome: str,
213    ) -> None:
214        log_event(
215            self._logger,
216            logging.DEBUG,
217            "provider_completed",
218            provider=provider,
219            stage="fetch",
220            url=target_url_for_log(str(url)),
221            outcome=outcome,
222            elapsed_ms=elapsed_ms(self._monotonic, started),
223        )
224
225    def _log_provider_failed(
226        self,
227        provider: str,
228        url: NormalizedURL,
229        started: float,
230        exc: Exception,
231        *,
232        failed_stage: str | None = None,
233        dependency: str | None = None,
234    ) -> None:
235        reason = exc.reason if isinstance(exc, ExecutionFailure) else None
236        if failed_stage is not None:
237            log_event(
238                self._logger,
239                logging.DEBUG,
240                "provider_failed",
241                provider=provider,
242                stage="fetch",
243                url=target_url_for_log(str(url)),
244                error_type=type(exc).__name__,
245                elapsed_ms=elapsed_ms(self._monotonic, started),
246                failed_stage=failed_stage,
247                dependency=dependency or "-",
248                reason=reason or "judge_execution_failed",
249            )
250            return
251        if reason:
252            log_event(
253                self._logger,
254                logging.DEBUG,
255                "provider_failed",
256                provider=provider,
257                stage="fetch",
258                url=target_url_for_log(str(url)),
259                error_type=type(exc).__name__,
260                elapsed_ms=elapsed_ms(self._monotonic, started),
261                reason=reason,
262            )
263            return
264        log_event(
265            self._logger,
266            logging.DEBUG,
267            "provider_failed",
268            provider=provider,
269            stage="fetch",
270            url=target_url_for_log(str(url)),
271            error_type=type(exc).__name__,
272            elapsed_ms=elapsed_ms(self._monotonic, started),
273        )
274
275    @staticmethod
276    def _validate_candidate(candidate: URLFetchCandidate) -> None:
277        if not isinstance(candidate, URLFetchCandidate):
278            raise TypeError("fetch candidate has invalid type")
279        if not isinstance(candidate.raw_content, str) or not isinstance(candidate.content, str):
280            raise TypeError("fetch candidate fields must be strings")
281        if candidate.raw_content == "":
282            raise ValueError("fetch candidate raw_content must be non-empty")