web_search.py

"""Web search functionality using Tavily API."""

from __future__ import annotations

from typing import List, Tuple
from urllib.parse import urlparse

from tavily import TavilyClient

from config import TAVILY_API_KEY, TRUSTED_DOMAINS
from helpers import _clip


def _domain(url: str) -> str:
    """Extract domain from URL."""
    netloc = urlparse(url).netloc.lower()
    if netloc.startswith("www."):
        netloc = netloc[4:]
    return netloc


def _is_legit_source(url: str) -> bool:
    """Check if URL is from a trusted domain."""
    dom = _domain(url)
    if not dom:
        return False
    if dom.endswith(".gov") or dom.endswith(".edu"):
        return True
    return any(dom == trusted or dom.endswith(f".{trusted}") for trusted in TRUSTED_DOMAINS)


def tavily_context(query: str, max_results: int = 6) -> Tuple[str, List[str], str]:
    """Search web using Tavily and return formatted context."""
    if not TAVILY_API_KEY:
        raise EnvironmentError("TAVILY_API_KEY is missing.")

    client = TavilyClient(api_key=TAVILY_API_KEY)
    response = client.search(
        query=query,
        max_results=max_results,
        search_depth="advanced",
        include_answer=False,
        include_raw_content=False,
    )
    results = response.get("results", [])
    if not results:
        return "", [], "Tavily returned no results."

    trusted = [item for item in results if _is_legit_source(item.get("url", ""))]
    selected = trusted[:5] if trusted else results[:5]
    status = (
        "Online search included trusted/official domains."
        if trusted
        else "No trusted-domain matches found; using top Tavily results."
    )

    snippets: List[str] = []
    source_refs: List[str] = []
    for idx, item in enumerate(selected, start=1):
        title = item.get("title", "Untitled")
        url = item.get("url", "")
        content = _clip(item.get("content", ""), max_chars=650)
        snippets.append(f"[WEB{idx}] {title} | {url}\n{content}")
        source_refs.append(f"- [WEB{idx}] {title} | {url}")

    return "\n\n".join(snippets), source_refs, status


def format_sources(report_sources: List[str], web_sources: List[str], web_status: str) -> str:
    """Format all source citations into a single string."""
    lines = ["Report citations:"]
    if report_sources:
        lines.extend(report_sources)
    else:
        lines.append("- None")

    if web_sources:
        lines.append("\nWeb citations:")
        lines.extend(web_sources)

    if web_status:
        lines.append(f"\nSearch note: {web_status}")

    return "\n".join(lines)