agent_search_gateway.orchestrators.fetch

  1"""URL fetch workflow orchestration and URL-state mutation."""
  2
  3import logging
  4import time
  5from collections.abc import Callable
  6
  7from ..concurrency import PerKeyLockPool, SingleflightGroup
  8from ..errors import UNAVAILABLE_MESSAGE, ErrorCode, ExecutionFailure, InputFailure
  9from ..llm.stages import LLMStages
 10from ..models import URLRecord
 11from ..observability import elapsed_ms, log_event, target_url_for_log
 12from ..scheduler.fetch import FetchScheduler
 13from ..url_normalization import NormalizedURL, normalize_url
 14from ..url_store import URLStore
 15
 16
 17class FetchOrchestrator:
 18    def __init__(
 19        self,
 20        *,
 21        store: URLStore,
 22        scheduler: FetchScheduler,
 23        stages: LLMStages,
 24        logger: logging.Logger | None = None,
 25        monotonic: Callable[[], float] = time.monotonic,
 26    ) -> None:
 27        self._store = store
 28        self._scheduler = scheduler
 29        self._stages = stages
 30        self._logger = logger or logging.getLogger(__name__)
 31        self._monotonic = monotonic
 32        self._url_locks: PerKeyLockPool[NormalizedURL] = PerKeyLockPool()
 33        self._request_singleflight: SingleflightGroup[tuple[NormalizedURL, str | None], str] = (
 34            SingleflightGroup()
 35        )
 36
 37    async def url_fetch(self, url: str, focus: str | None = None) -> str:
 38        normalized_url = normalize_url(url)
 39        normalized_focus = focus.strip() if focus is not None and focus.strip() else None
 40        key = (normalized_url, normalized_focus)
 41        focus_present = normalized_focus is not None
 42        focus_chars = len(normalized_focus or "")
 43        return await self._request_singleflight.do(
 44            key,
 45            lambda: self._serialized_url_fetch(normalized_url, normalized_focus),
 46            on_leader=lambda: log_event(
 47                self._logger,
 48                logging.DEBUG,
 49                "singleflight_leader",
 50                url=target_url_for_log(str(normalized_url)),
 51                focus_present=focus_present,
 52                focus_chars=focus_chars,
 53            ),
 54            on_follower=lambda: log_event(
 55                self._logger,
 56                logging.DEBUG,
 57                "singleflight_joined",
 58                url=target_url_for_log(str(normalized_url)),
 59                focus_present=focus_present,
 60                focus_chars=focus_chars,
 61            ),
 62        )
 63
 64    async def _serialized_url_fetch(
 65        self,
 66        normalized_url: NormalizedURL,
 67        normalized_focus: str | None,
 68    ) -> str:
 69        lock_started = self._monotonic()
 70        async with self._url_locks.acquire(normalized_url):
 71            log_event(
 72                self._logger,
 73                logging.DEBUG,
 74                "url_lock_acquired",
 75                url=target_url_for_log(str(normalized_url)),
 76                wait_ms=elapsed_ms(self._monotonic, lock_started),
 77            )
 78            snapshot = self._store.get(normalized_url)
 79            if snapshot is None:
 80                self._log_state(normalized_url, "url_not_admitted")
 81                raise InputFailure(ErrorCode.URL_NOT_ADMITTED, "URL was not admitted by search")
 82            if not snapshot.available:
 83                self._log_state(normalized_url, "stored_unavailable")
 84                return UNAVAILABLE_MESSAGE
 85
 86            prepared = await self._prepare_content(normalized_url)
 87            if not prepared:
 88                return UNAVAILABLE_MESSAGE
 89
 90            refreshed = self._require_snapshot(normalized_url)
 91            if not await self._safety_check(normalized_url, refreshed.content):
 92                return UNAVAILABLE_MESSAGE
 93            if normalized_focus is None:
 94                self._log_final(normalized_url, refreshed.content, "content")
 95                return refreshed.content
 96            summary = await self._focus_summary(refreshed.content, normalized_focus)
 97            self._log_final(normalized_url, summary, "focus_summary")
 98            return summary
 99
100    async def _prepare_content(self, url: NormalizedURL) -> bool:
101        snapshot = self._require_snapshot(url)
102        if snapshot.content:
103            self._log_state(url, "cached_content")
104            return True
105
106        if not snapshot.raw_content:
107            self._log_state(url, "provider_fetch_required")
108            if not self._scheduler.provider_names:
109                raise ExecutionFailure(
110                    ErrorCode.NO_URL_FETCH_PROVIDERS,
111                    "No URL fetch providers are enabled",
112                )
113            outcome = await self._scheduler.fetch_until_accepted(url)
114            if outcome.kind == "execution_failure":
115                raise ExecutionFailure(
116                    ErrorCode.ALL_PROVIDERS_FAILED,
117                    "All URL fetch provider pipelines failed",
118                )
119            if outcome.kind == "semantic_failure":
120                self._store.mark_unavailable(url)
121                self._log_rejected(url, "fetch_semantic_failure")
122                return False
123            candidate = outcome.candidate
124            if candidate is None:
125                raise RuntimeError("accepted fetch outcome did not contain a candidate")
126            self._store.merge_body(
127                url,
128                raw_content=candidate.raw_content,
129                content=candidate.content,
130            )
131            log_event(
132                self._logger,
133                logging.DEBUG,
134                "body_accepted",
135                url=target_url_for_log(str(url)),
136                reason="provider_fetch",
137                raw_chars=len(candidate.raw_content),
138                content_chars=len(candidate.content),
139            )
140            snapshot = self._require_snapshot(url)
141        else:
142            self._log_state(url, "raw_content_available")
143
144        if not snapshot.content:
145            cleaned = await self._content_clean(snapshot.raw_content)
146            self._store.merge_body(url, content=cleaned)
147            log_event(
148                self._logger,
149                logging.DEBUG,
150                "body_accepted",
151                url=target_url_for_log(str(url)),
152                reason="content_cleaned",
153                content_chars=len(cleaned),
154            )
155        return True
156
157    async def _content_clean(self, raw_content: str) -> str:
158        try:
159            return await self._stages.content_clean(raw_content)
160        except ExecutionFailure as exc:
161            raise ExecutionFailure(
162                ErrorCode.LLM_STAGE_FAILED,
163                "Content-clean LLM stage failed",
164            ) from exc
165
166    async def _safety_check(self, url: NormalizedURL, content: str) -> bool:
167        try:
168            decision = await self._stages.safety(content)
169        except ExecutionFailure as exc:
170            raise ExecutionFailure(
171                ErrorCode.LLM_STAGE_FAILED,
172                "Safety LLM stage failed",
173            ) from exc
174        if decision.ok:
175            return True
176        self._store.mark_unavailable(url)
177        self._log_rejected(url, "safety_rejected")
178        return False
179
180    async def _focus_summary(self, content: str, focus: str) -> str:
181        try:
182            return await self._stages.focus_summary(content, focus)
183        except ExecutionFailure as exc:
184            raise ExecutionFailure(
185                ErrorCode.LLM_STAGE_FAILED,
186                "Focus-summary LLM stage failed",
187            ) from exc
188
189    def _log_state(self, url: NormalizedURL, reason: str) -> None:
190        log_event(
191            self._logger,
192            logging.DEBUG,
193            "body_skipped",
194            url=target_url_for_log(str(url)),
195            reason=reason,
196        )
197
198    def _log_rejected(self, url: NormalizedURL, reason: str) -> None:
199        log_event(
200            self._logger,
201            logging.DEBUG,
202            "body_rejected",
203            url=target_url_for_log(str(url)),
204            reason=reason,
205        )
206
207    def _log_final(self, url: NormalizedURL, text: str, reason: str) -> None:
208        log_event(
209            self._logger,
210            logging.DEBUG,
211            "body_accepted",
212            url=target_url_for_log(str(url)),
213            reason=reason,
214            output_chars=len(text),
215        )
216
217    def _require_snapshot(self, url: NormalizedURL) -> URLRecord:
218        snapshot = self._store.get(url)
219        if snapshot is None:
220            raise RuntimeError("admitted URL disappeared from store")
221        return snapshot