from __future__ import annotations import asyncio import os from datetime import datetime, timezone from typing import Any, Awaitable, Callable from ..schemas import ExternalAdapterResult, ExternalEvidenceItem, ExternalQueryPlan class LiveSearchFallbackAdapter: adapter_name = "live_search" supported_query_kinds = {"weather", "news", "market", "generic_web", "unknown"} def __init__( self, *, build_queries_callable: Callable[[str], list[str]] | None = None, search_results_callable: Callable[..., Awaitable[list[dict[str, str]]]] | None = None, fetch_text_callable: Callable[[str], Awaitable[str]] | None = None, chunk_text_callable: Callable[..., list[str]] | None = None, score_search_result_callable: Callable[..., float] | None = None, score_chunk_callable: Callable[..., float] | None = None, max_pages: int = 3, page_fetch_timeout_seconds: float | None = None, ) -> None: self.build_queries_callable = build_queries_callable self.search_results_callable = search_results_callable self.fetch_text_callable = fetch_text_callable self.chunk_text_callable = chunk_text_callable self.score_search_result_callable = score_search_result_callable self.score_chunk_callable = score_chunk_callable self.max_pages = max(1, int(max_pages or 3)) self.page_fetch_timeout_seconds = _normalize_page_fetch_timeout(page_fetch_timeout_seconds) async def execute(self, plan: ExternalQueryPlan) -> ExternalAdapterResult: if not all( callable(item) for item in ( self.search_results_callable, self.fetch_text_callable, self.chunk_text_callable, self.score_search_result_callable, self.score_chunk_callable, ) ): return ExternalAdapterResult( adapter_name=self.adapter_name, status="unavailable", evidence=[], meta={"reason": "live_search_callbacks_missing"}, ) queries = [str(item or "").strip() for item in list(plan.search_queries or []) if str(item or "").strip()] if not queries and callable(self.build_queries_callable): try: queries = [ str(item or "").strip() for item in list(self.build_queries_callable(plan.topic or "") or []) if str(item or "").strip() ] except Exception: queries = [] if not queries: return ExternalAdapterResult(adapter_name=self.adapter_name, status="empty", evidence=[], meta={"query": ""}) seen_urls: set[str] = set() query_limit = min(len(queries), 5) if (bool(plan.freshness_required) or bool(plan.requires_multi_source)) else len(queries) active_queries = queries[:query_limit] async def _search(query: str) -> tuple[str, list[dict[str, str]]]: try: rows = await self.search_results_callable( query=query, max_results=max(3, int(plan.result_limit) * 2), ) except Exception: rows = [] return query, [dict(row) for row in list(rows or []) if isinstance(row, dict)] search_batches = await asyncio.gather(*[_search(query) for query in active_queries], return_exceptions=True) search_results: list[dict[str, str]] = [] selected_query = "" for batch in search_batches: if isinstance(batch, Exception): continue query, rows = batch if rows and not selected_query: selected_query = query for row in rows: url = str(row.get("url") or "").strip() if not url or url in seen_urls: continue seen_urls.add(url) copied = dict(row) copied["_query"] = query search_results.append(copied) ranked = sorted( search_results, key=lambda row: float( self.score_search_result_callable( question=plan.topic or selected_query or queries[0], title=str(row.get("title") or ""), url=str(row.get("url") or ""), ) or 0.0 ), reverse=True, ) if bool(plan.requires_multi_source): ranked = _diversify_ranked_results_by_query(ranked, limit=max(self.max_pages + len(queries), int(plan.result_limit) * 2)) candidate_limit = min( len(ranked), max(self.max_pages + 2, int(plan.result_limit) * 2) if bool(plan.requires_multi_source) else self.max_pages + 2, ) async def _fetch_evidence(row: dict[str, str]) -> tuple[ExternalEvidenceItem | None, bool]: url = str(row.get("url") or "").strip() title = str(row.get("title") or "").strip() row_query = str(row.get("_query") or selected_query or queries[0]).strip() if not url: return None, False try: text = await asyncio.wait_for( self.fetch_text_callable(url), timeout=self.page_fetch_timeout_seconds, ) except Exception: return None, True chunks = list(self.chunk_text_callable(text, chunk_size=900, overlap=120) or []) if not chunks: return None, True best_chunk = "" best_score = 0.0 for chunk in chunks[:6]: score = float( self.score_chunk_callable( question=plan.topic or selected_query or queries[0], title=title, content=chunk, url=url, ) or 0.0 ) if score > best_score: best_score = score best_chunk = str(chunk or "").strip() if not best_chunk: return None, True return ( ExternalEvidenceItem( source_type="url", source_label=title or url, title=title[:255], target_url=url, snippet=best_chunk[:1800], retrieved_at=datetime.now(timezone.utc).isoformat(), freshness_score=round(best_score, 4), authority_score=round(best_score, 4), grounded_fields=["snippet"], raw_metadata={ "query": row_query, "result_query": row_query, "provider": self.adapter_name, "title": title[:255], }, ), True, ) fetch_rows = ranked[:candidate_limit] evidence: list[ExternalEvidenceItem] = [] fetch_attempts = 0 cancelled_fetches = 0 fetch_order = {str(row.get("url") or "").strip(): index for index, row in enumerate(fetch_rows)} pending = {asyncio.create_task(_fetch_evidence(row)) for row in fetch_rows} evidence_target = max(1, int(plan.result_limit)) loop = asyncio.get_running_loop() deadline = loop.time() + self.page_fetch_timeout_seconds while pending and len(evidence) < evidence_target: remaining = max(0.0, deadline - loop.time()) if remaining <= 0: break done, pending = await asyncio.wait( pending, timeout=remaining, return_when=asyncio.FIRST_COMPLETED, ) if not done: break for task in done: try: evidence_item, attempted = task.result() except Exception: continue if attempted: fetch_attempts += 1 if evidence_item is not None: evidence.append(evidence_item) if pending: cancelled_fetches = len(pending) for task in pending: task.cancel() evidence.sort(key=lambda item: fetch_order.get(str(item.target_url or "").strip(), 10_000)) return ExternalAdapterResult( adapter_name=self.adapter_name, status="success" if evidence else "empty", evidence=evidence[: max(1, int(plan.result_limit))], meta={ "query": selected_query or (queries[0] if queries else ""), "search_result_count": len(search_results), "fetched_pages": fetch_attempts, "cancelled_fetches": cancelled_fetches, "page_fetch_timeout_seconds": self.page_fetch_timeout_seconds, }, ) def _diversify_ranked_results_by_query(rows: list[dict[str, str]], *, limit: int) -> list[dict[str, str]]: """Prefer coverage across planner-provided query lanes before taking extra same-lane hits.""" if not rows: return [] grouped: dict[str, list[dict[str, str]]] = {} query_order: list[str] = [] for row in rows: query = str(row.get("_query") or "").strip() if query not in grouped: grouped[query] = [] query_order.append(query) grouped[query].append(row) diversified: list[dict[str, str]] = [] seen_urls: set[str] = set() while len(diversified) < max(1, int(limit or 1)): changed = False for query in query_order: bucket = grouped.get(query) or [] while bucket: candidate = bucket.pop(0) url = str(candidate.get("url") or "").strip() if url and url in seen_urls: continue if url: seen_urls.add(url) diversified.append(candidate) changed = True break if len(diversified) >= max(1, int(limit or 1)): break if not changed: break return diversified def _normalize_page_fetch_timeout(value: float | None) -> float: if value is None: raw = os.getenv("AI_EXTERNAL_LIVE_PAGE_FETCH_TIMEOUT_SECONDS", "8") try: value = float(raw) except Exception: value = 8.0 try: timeout = float(value) except Exception: timeout = 8.0 return min(max(timeout, 3.0), 12.0)