agent_search_gateway.orchestrators.search

  1"""Keyword and LLM search workflow orchestration."""
  2
  3import asyncio
  4import logging
  5import time
  6from collections.abc import Callable, Sequence
  7from dataclasses import dataclass
  8
  9from ..academic.aggregator import PaperAggregator
 10from ..concurrency import ProviderQuotaManager
 11from ..errors import ErrorCode, ExecutionFailure, InputFailure
 12from ..llm.stages import LLMStages, cheap_check
 13from ..llm_search_parser import parse_search_markdown
 14from ..models import LLMInvocation, LLMSearchScope, PaperRecord, SearchRecord
 15from ..observability import elapsed_ms, log_event, target_url_for_log
 16from ..paper_search_parser import parse_paper_markdown
 17from ..providers.contracts import (
 18    KeywordSearchHit,
 19    KeywordSearchProvider,
 20    OAResolver,
 21    PaperSearchHit,
 22)
 23from ..request_ids import validate_request_id
 24from ..result_writer import ResultWriter
 25from ..url_normalization import NormalizedURL, normalize_url
 26from ..url_store import URLStore
 27from .paper import finalize_paper_hits
 28
 29
 30@dataclass(frozen=True, slots=True)
 31class _StagedKeyword:
 32    url: NormalizedURL
 33    abstract: str
 34    provider: str
 35    raw_content: str = ""
 36    content: str = ""
 37
 38
 39class SearchOrchestrator:
 40    def __init__(
 41        self,
 42        *,
 43        keyword_providers: Sequence[KeywordSearchProvider],
 44        llm_invocations: Sequence[LLMInvocation],
 45        quotas: ProviderQuotaManager,
 46        stages: LLMStages,
 47        store: URLStore,
 48        result_writer: ResultWriter,
 49        paper_aggregator: PaperAggregator | None = None,
 50        paper_resolver: OAResolver | None = None,
 51        logger: logging.Logger | None = None,
 52        monotonic: Callable[[], float] = time.monotonic,
 53    ) -> None:
 54        self._keyword_providers = tuple(keyword_providers)
 55        self._llm_invocations = tuple(llm_invocations)
 56        self.quotas = quotas
 57        self._stages = stages
 58        self._store = store
 59        self._result_writer = result_writer
 60        self._paper_aggregator = paper_aggregator or PaperAggregator(
 61            tuple(f"llm:{invocation.provider}" for invocation in self._llm_invocations)
 62        )
 63        self._paper_resolver = paper_resolver
 64        self._logger = logger or logging.getLogger(__name__)
 65        self._monotonic = monotonic
 66
 67    async def keyword_search(self, query: str, *, request_id: str) -> str:
 68        validate_request_id(request_id)
 69        normalized_query = query.strip()
 70        if not normalized_query:
 71            raise InputFailure(ErrorCode.EMPTY_QUERY, "Query must not be empty")
 72        if not self._keyword_providers:
 73            raise ExecutionFailure(
 74                ErrorCode.NO_KEYWORD_SEARCH_PROVIDERS,
 75                "No keyword search providers are enabled",
 76            )
 77
 78        outcomes = await asyncio.gather(
 79            *(
 80                self._run_keyword_pipeline(provider, normalized_query)
 81                for provider in self._keyword_providers
 82            ),
 83            return_exceptions=True,
 84        )
 85        completed = [outcome for outcome in outcomes if isinstance(outcome, list)]
 86        if not completed:
 87            raise ExecutionFailure(
 88                ErrorCode.ALL_PROVIDERS_FAILED,
 89                "All keyword search provider pipelines failed",
 90            )
 91
 92        ordered_urls: list[NormalizedURL] = []
 93        seen: set[NormalizedURL] = set()
 94        for outcome in outcomes:
 95            if not isinstance(outcome, list):
 96                continue
 97            for staged in outcome:
 98                self._store.admit(
 99                    staged.url,
100                    staged.abstract,
101                    raw_content=staged.raw_content,
102                    content=staged.content,
103                )
104                if staged.url not in seen:
105                    seen.add(staged.url)
106                    ordered_urls.append(staged.url)
107                else:
108                    log_event(
109                        self._logger,
110                        logging.DEBUG,
111                        "candidate_rejected",
112                        provider=staged.provider,
113                        url=target_url_for_log(str(staged.url)),
114                        reason="duplicate",
115                    )
116
117        records = [self._record_from_store(url) for url in ordered_urls]
118        path = self._result_writer.write_results("keyword", records, request_id=request_id)
119        log_event(
120            self._logger,
121            logging.DEBUG,
122            "results_written",
123            kind="keyword",
124            path=str(path),
125            results=len(records),
126        )
127        return str(path)
128
129    async def llm_search(
130        self,
131        prompt: str,
132        *,
133        request_id: str,
134        scope: LLMSearchScope = "web",
135    ) -> str:
136        validate_request_id(request_id)
137        normalized_prompt = prompt.strip()
138        if not normalized_prompt:
139            raise InputFailure(ErrorCode.EMPTY_QUERY, "Prompt must not be empty")
140        if scope not in {"web", "paper", "all"}:
141            raise InputFailure(ErrorCode.BAD_REQUEST, "LLM search scope is invalid")
142        if not self._llm_invocations:
143            raise ExecutionFailure(
144                ErrorCode.NO_LLM_SEARCH_PROVIDERS,
145                "No LLM search providers are configured",
146            )
147
148        if scope == "web":
149            web_records = await self._llm_web_records(normalized_prompt)
150            path = self._result_writer.write_results("llm", web_records, request_id=request_id)
151            self._log_results_written(path=str(path), results=len(web_records))
152            return str(path)
153
154        if scope == "paper":
155            paper_records = await self._llm_paper_records(normalized_prompt)
156            path = self._result_writer.write_paper_results(
157                "llm",
158                paper_records,
159                request_id=request_id,
160            )
161            self._log_results_written(path=str(path), results=len(paper_records))
162            return str(path)
163
164        web_outcome, paper_outcome = await asyncio.gather(
165            self._llm_web_records(normalized_prompt),
166            self._llm_paper_records(normalized_prompt),
167            return_exceptions=True,
168        )
169        if isinstance(web_outcome, BaseException):
170            self._log_branch_failure("web", web_outcome)
171            mixed_web_records: list[SearchRecord] = []
172        else:
173            mixed_web_records = web_outcome
174        if isinstance(paper_outcome, BaseException):
175            self._log_branch_failure("paper", paper_outcome)
176            mixed_paper_records: list[PaperRecord] = []
177        else:
178            mixed_paper_records = paper_outcome
179        if isinstance(web_outcome, BaseException) and isinstance(
180            paper_outcome,
181            BaseException,
182        ):
183            raise ExecutionFailure(
184                ErrorCode.ALL_PROVIDERS_FAILED,
185                "All LLM search branches failed",
186            )
187        path = self._result_writer.write_mixed_results(
188            mixed_web_records,
189            mixed_paper_records,
190            request_id=request_id,
191        )
192        self._log_results_written(
193            path=str(path),
194            results=len(mixed_web_records) + len(mixed_paper_records),
195        )
196        return str(path)
197
198    async def _llm_web_records(self, prompt: str) -> list[SearchRecord]:
199        outcomes = await asyncio.gather(
200            *(self._run_llm_pipeline(invocation, prompt) for invocation in self._llm_invocations),
201            return_exceptions=True,
202        )
203        if not any(isinstance(outcome, list) for outcome in outcomes):
204            raise ExecutionFailure(
205                ErrorCode.ALL_PROVIDERS_FAILED,
206                "All LLM web search provider pipelines failed",
207            )
208        ordered_urls: list[NormalizedURL] = []
209        seen: set[NormalizedURL] = set()
210        for outcome in outcomes:
211            if not isinstance(outcome, list):
212                continue
213            for result in outcome:
214                self._store.admit(result.url, result.abstract)
215                if result.url not in seen:
216                    seen.add(result.url)
217                    ordered_urls.append(result.url)
218        return [self._record_from_store(url) for url in ordered_urls]
219
220    async def _llm_paper_records(self, prompt: str) -> list[PaperRecord]:
221        outcomes = await asyncio.gather(
222            *(
223                self._run_llm_paper_pipeline(invocation, prompt)
224                for invocation in self._llm_invocations
225            ),
226            return_exceptions=True,
227        )
228        if not any(isinstance(outcome, list) for outcome in outcomes):
229            raise ExecutionFailure(
230                ErrorCode.ALL_PROVIDERS_FAILED,
231                "All LLM paper search provider pipelines failed",
232            )
233        hits: list[PaperSearchHit] = []
234        for outcome in outcomes:
235            if isinstance(outcome, list):
236                hits.extend(outcome)
237        return await finalize_paper_hits(
238            hits,
239            aggregator=self._paper_aggregator,
240            resolver=self._paper_resolver,
241            store=self._store,
242        )
243
244    def _log_results_written(self, *, path: str, results: int) -> None:
245        log_event(
246            self._logger,
247            logging.DEBUG,
248            "results_written",
249            kind="llm",
250            path=path,
251            results=results,
252        )
253
254    def _log_branch_failure(self, scope: str, exc: BaseException) -> None:
255        log_event(
256            self._logger,
257            logging.DEBUG,
258            "llm_search_branch_failed",
259            scope=scope,
260            error_type=type(exc).__name__,
261        )
262
263    async def _run_llm_pipeline(
264        self,
265        invocation: LLMInvocation,
266        prompt: str,
267    ) -> list[SearchRecord]:
268        started = self._monotonic()
269        log_event(
270            self._logger,
271            logging.DEBUG,
272            "provider_started",
273            provider=invocation.provider,
274            stage="llm_search",
275            model=invocation.model,
276        )
277        try:
278            markdown = await self._stages.llm_search_markdown(invocation, prompt)
279            records = parse_search_markdown(markdown)
280        except asyncio.CancelledError:
281            raise
282        except ExecutionFailure as exc:
283            self._log_provider_failure(invocation.provider, "llm_search", started, exc)
284            raise
285        except Exception as exc:
286            self._log_provider_failure(invocation.provider, "llm_search", started, exc)
287            raise ExecutionFailure(
288                ErrorCode.ALL_PROVIDERS_FAILED,
289                f"LLM search provider {invocation.provider} returned invalid data",
290            ) from exc
291        log_event(
292            self._logger,
293            logging.DEBUG,
294            "provider_completed",
295            provider=invocation.provider,
296            stage="llm_search",
297            model=invocation.model,
298            output_chars=len(markdown),
299            results=len(records),
300            elapsed_ms=elapsed_ms(self._monotonic, started),
301        )
302        return records
303
304    async def _run_llm_paper_pipeline(
305        self,
306        invocation: LLMInvocation,
307        prompt: str,
308    ) -> list[PaperSearchHit]:
309        started = self._monotonic()
310        log_event(
311            self._logger,
312            logging.DEBUG,
313            "provider_started",
314            provider=invocation.provider,
315            stage="llm_paper_search",
316            model=invocation.model,
317        )
318        try:
319            markdown = await self._stages.llm_paper_search_markdown(invocation, prompt)
320            hits = parse_paper_markdown(markdown, provider=invocation.provider)
321        except asyncio.CancelledError:
322            raise
323        except ExecutionFailure as exc:
324            self._log_provider_failure(invocation.provider, "llm_paper_search", started, exc)
325            raise
326        except Exception as exc:
327            self._log_provider_failure(invocation.provider, "llm_paper_search", started, exc)
328            raise ExecutionFailure(
329                ErrorCode.ALL_PROVIDERS_FAILED,
330                f"LLM paper search provider {invocation.provider} returned invalid data",
331            ) from exc
332        log_event(
333            self._logger,
334            logging.DEBUG,
335            "provider_completed",
336            provider=invocation.provider,
337            stage="llm_paper_search",
338            model=invocation.model,
339            output_chars=len(markdown),
340            results=len(hits),
341            elapsed_ms=elapsed_ms(self._monotonic, started),
342        )
343        return hits
344
345    async def _run_keyword_pipeline(
346        self,
347        provider: KeywordSearchProvider,
348        query: str,
349    ) -> list[_StagedKeyword]:
350        started = self._monotonic()
351        log_event(
352            self._logger,
353            logging.DEBUG,
354            "provider_started",
355            provider=provider.name,
356            stage="search",
357        )
358        try:
359            async with self.quotas.get_web(provider.name).lease():
360                hits = await provider.search(query)
361            if not isinstance(hits, list):
362                raise TypeError("provider search result must be a list")
363        except asyncio.CancelledError:
364            raise
365        except ExecutionFailure as exc:
366            self._log_provider_failure(provider.name, "search", started, exc)
367            raise
368        except Exception as exc:
369            self._log_provider_failure(provider.name, "search", started, exc)
370            raise ExecutionFailure(
371                ErrorCode.ALL_PROVIDERS_FAILED,
372                f"Keyword provider {provider.name} returned invalid data",
373            ) from exc
374
375        staged: list[_StagedKeyword] = []
376        hit_failure: ExecutionFailure | None = None
377        failed_hits = 0
378        try:
379            for hit in hits:
380                try:
381                    staged_hit = await self._stage_keyword_hit(hit, provider=provider.name)
382                except ExecutionFailure as exc:
383                    # Keyword hits are independent; one failed judge must not discard sibling hits.
384                    if hit_failure is None:
385                        hit_failure = exc
386                    failed_hits += 1
387                    continue
388                if staged_hit is not None:
389                    staged.append(staged_hit)
390        except asyncio.CancelledError:
391            raise
392        except Exception as exc:
393            self._log_provider_failure(provider.name, "hit_staging", started, exc)
394            raise ExecutionFailure(
395                ErrorCode.ALL_PROVIDERS_FAILED,
396                f"Keyword provider {provider.name} returned invalid data",
397            ) from exc
398
399        if hit_failure is not None:
400            if not staged:
401                self._log_provider_failure(provider.name, "judge", started, hit_failure)
402                raise hit_failure
403            self._log_provider_partial_failure(
404                provider.name,
405                "judge",
406                started,
407                hit_failure,
408                failed_hits=failed_hits,
409            )
410        log_event(
411            self._logger,
412            logging.DEBUG,
413            "provider_completed",
414            provider=provider.name,
415            stage="search",
416            hits=len(hits),
417            results=len(staged),
418            elapsed_ms=elapsed_ms(self._monotonic, started),
419        )
420        return staged
421
422    async def _stage_keyword_hit(
423        self,
424        hit: KeywordSearchHit,
425        *,
426        provider: str,
427    ) -> _StagedKeyword | None:
428        if not isinstance(hit, KeywordSearchHit):
429            raise TypeError("keyword hit has invalid type")
430        for value in (hit.url, hit.title, hit.snippet, hit.raw_content, hit.content):
431            if not isinstance(value, str):
432                raise TypeError("keyword hit fields must be strings")
433
434        abstract = hit.snippet.strip() or hit.title.strip()
435        if not abstract:
436            try:
437                logged_url = target_url_for_log(str(normalize_url(hit.url)))
438            except InputFailure:
439                logged_url = target_url_for_log(hit.url)
440            log_event(
441                self._logger,
442                logging.DEBUG,
443                "candidate_rejected",
444                provider=provider,
445                url=logged_url,
446                reason="empty_abstract",
447            )
448            return None
449        url = normalize_url(hit.url)
450        log_event(
451            self._logger,
452            logging.DEBUG,
453            "candidate_accepted",
454            provider=provider,
455            url=target_url_for_log(str(url)),
456            abstract_chars=len(abstract),
457        )
458        current = self._store.get(url)
459        if current is not None and not current.available:
460            self._log_body_decision(provider, url, "body_skipped", "stored_unavailable")
461            return _StagedKeyword(url=url, abstract=abstract, provider=provider)
462
463        had_body = bool(hit.raw_content or hit.content)
464        raw_content = hit.raw_content if hit.raw_content.strip() else ""
465        content = hit.content if hit.content.strip() else ""
466        candidate = content or raw_content
467        if not candidate:
468            reason = "cheap_check" if had_body else "no_body"
469            event = "body_rejected" if had_body else "body_skipped"
470            self._log_body_decision(provider, url, event, reason)
471            return _StagedKeyword(url=url, abstract=abstract, provider=provider)
472        if not cheap_check(candidate):
473            self._log_body_decision(provider, url, "body_rejected", "cheap_check")
474            return _StagedKeyword(url=url, abstract=abstract, provider=provider)
475
476        decision = await self._stages.judge(candidate)
477        if not decision.ok:
478            self._log_body_decision(provider, url, "body_rejected", "judge_rejected")
479            return _StagedKeyword(url=url, abstract=abstract, provider=provider)
480        log_event(
481            self._logger,
482            logging.DEBUG,
483            "body_accepted",
484            provider=provider,
485            url=target_url_for_log(str(url)),
486            raw_chars=len(raw_content),
487            content_chars=len(content),
488        )
489        return _StagedKeyword(
490            url=url,
491            abstract=abstract,
492            provider=provider,
493            raw_content=raw_content,
494            content=content,
495        )
496
497    def _log_body_decision(
498        self,
499        provider: str,
500        url: NormalizedURL,
501        event: str,
502        reason: str,
503    ) -> None:
504        log_event(
505            self._logger,
506            logging.DEBUG,
507            event,
508            provider=provider,
509            url=target_url_for_log(str(url)),
510            reason=reason,
511        )
512
513    def _log_provider_failure(
514        self,
515        provider: str,
516        stage: str,
517        started: float,
518        exc: Exception,
519    ) -> None:
520        log_event(
521            self._logger,
522            logging.DEBUG,
523            "provider_failed",
524            provider=provider,
525            stage=stage,
526            error_type=type(exc).__name__,
527            elapsed_ms=elapsed_ms(self._monotonic, started),
528        )
529
530    def _log_provider_partial_failure(
531        self,
532        provider: str,
533        stage: str,
534        started: float,
535        exc: Exception,
536        *,
537        failed_hits: int,
538    ) -> None:
539        log_event(
540            self._logger,
541            logging.DEBUG,
542            "provider_partial_failure",
543            provider=provider,
544            stage=stage,
545            error_type=type(exc).__name__,
546            failed_hits=failed_hits,
547            elapsed_ms=elapsed_ms(self._monotonic, started),
548        )
549
550    def _record_from_store(self, url: NormalizedURL) -> SearchRecord:
551        record = self._store.get(url)
552        if record is None:
553            raise RuntimeError("committed URL disappeared from store")
554        return SearchRecord(record.url, record.abstract)