PMFlow/backend/app/services/external_query_service/adapters/live_search.py

from __future__ import annotations

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,
    ) -> 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))

    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": ""})

        search_results: list[dict[str, str]] = []
        selected_query = ""
        for query in queries:
            try:
                search_results = await self.search_results_callable(
                    query=query,
                    max_results=max(3, int(plan.result_limit) * 2),
                )
            except Exception:
                search_results = []
            if search_results:
                selected_query = query
                break

        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,
        )

        evidence: list[ExternalEvidenceItem] = []
        fetch_attempts = 0
        candidate_limit = min(len(ranked), self.max_pages + 2)
        for row in ranked[:candidate_limit]:
            url = str(row.get("url") or "").strip()
            title = str(row.get("title") or "").strip()
            if not url:
                continue
            fetch_attempts += 1
            try:
                text = await self.fetch_text_callable(url)
            except Exception:
                continue
            chunks = list(self.chunk_text_callable(text, chunk_size=900, overlap=120) or [])
            if not chunks:
                continue
            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:
                continue
            evidence.append(
                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": selected_query or queries[0],
                        "provider": self.adapter_name,
                        "title": title[:255],
                    },
                )
            )

        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,
            },
        )