272 lines
11 KiB
Python
272 lines
11 KiB
Python
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)
|