Axon/axon/routers/search.py

import asyncio
import json
import logging
import time
from urllib.parse import urlparse

from fastapi import APIRouter, HTTPException

from axon.models import SearchRequest, SearchResponse, SearchResult
from axon.services.cache import get_cache, make_cache_key
from axon.services.index import search_index_as_results
from axon.services.quality import annotate_source_quality
from axon.services.reranker import rerank_results
from axon.services.scraper import scrape_urls
from axon.services.sources.brave import search_brave
from axon.services.sources.ddg import search_ddg
from axon.services.sources.serper import search_serper
from axon.services.synthesizer import synthesize_answer
from axon.config import settings

router = APIRouter()
logger = logging.getLogger(__name__)


def deduplicate_results(results: list[SearchResult]) -> list[SearchResult]:
    """Keep first occurrence of each URL, preserve order."""
    seen: set[str] = set()
    out: list[SearchResult] = []
    for r in results:
        normalized = r.url.rstrip("/").lower()
        if normalized not in seen:
            seen.add(normalized)
            out.append(r)
    return out


def _filter(
    results: list[SearchResult],
    whitelist: list[str],
    blacklist: list[str],
) -> list[SearchResult]:
    if whitelist:
        allowed = {d.lower().removeprefix("www.") for d in whitelist}
        results = [
            r for r in results
            if urlparse(r.url).hostname
            and urlparse(r.url).hostname.lower().removeprefix("www.") in allowed
        ]
    if blacklist:
        blocked = {d.lower().removeprefix("www.") for d in blacklist}
        results = [
            r for r in results
            if not (
                urlparse(r.url).hostname
                and urlparse(r.url).hostname.lower().removeprefix("www.") in blocked
            )
        ]
    return results


@router.post("/search", response_model=SearchResponse)
async def search(request: SearchRequest) -> SearchResponse:
    start = time.monotonic()
    cache = await get_cache()
    cache_key = make_cache_key(
        "search",
        query=request.query,
        max_results=request.max_results,
        search_depth=request.search_depth,
        include_answer=request.include_answer,
        include_raw_content=request.include_raw_content,
        include_private=request.include_private,
        topic=request.topic,
        time_range=request.time_range,
        whitelist=sorted(request.domains_whitelist),
        blacklist=sorted(request.domains_blacklist),
    )

    cached = await cache.get(cache_key)
    if cached:
        data = json.loads(cached)
        data["response_time"] = round(time.monotonic() - start, 3)
        return SearchResponse(**data)

    # Fetch sources concurrently; Serper is conditional
    brave_res, ddg_res, private_res = await asyncio.gather(
        search_brave(request.query, count=request.max_results, topic=request.topic, time_range=request.time_range),
        search_ddg(request.query, count=request.max_results, topic=request.topic, time_range=request.time_range),
        search_index_as_results(request.query, max_results=min(5, request.max_results))
        if request.include_private else asyncio.sleep(0, result=[]),
    )

    serper_res: list[SearchResult] = []
    if len(brave_res) < 5:
        serper_res = await search_serper(
            request.query, count=request.max_results, topic=request.topic, time_range=request.time_range
        )

    all_results = private_res + brave_res + ddg_res + serper_res
    if not all_results:
        raise HTTPException(status_code=503, detail="All search sources returned no results")

    results = deduplicate_results(all_results)
    results = _filter(results, request.domains_whitelist, request.domains_blacklist)
    results = results[: request.max_results]

    source_count = len(results)
    scrape_count = 0

    if request.search_depth == "advanced" and results:
        scraped = await scrape_urls(
            [r.url for r in results],
            max_concurrent=settings.max_concurrent_scrapes,
            timeout=settings.scrape_timeout,
        )
        scrape_count = sum(1 for v in scraped.values() if v)
        contents = [(r.url, scraped.get(r.url) or r.content) for r in results]
    else:
        contents = [(r.url, r.content) for r in results]

    ranked = await rerank_results(request.query, contents)
    url_rank: dict[str, tuple[str, float]] = {url: (chunk, score) for url, chunk, score in ranked}

    for r in results:
        if r.url in url_rank:
            chunk, score = url_rank[r.url]
            r.content = chunk or r.content
            r.score = score

    results.sort(key=lambda r: r.score, reverse=True)
    annotate_source_quality(results)

    if not request.include_raw_content:
        for r in results:
            r.content = r.content[:500]

    answer: str | None = None
    if request.include_answer:
        answer = await synthesize_answer(request.query, results)

    scores = [r.score for r in results]
    response_time = round(time.monotonic() - start, 3)
    logger.info(
        "search query=%r private=%d brave=%d ddg=%d serper=%d returned=%d scrapes=%d scores=[%.3f,%.3f] time=%.3fs",
        request.query,
        len(private_res),
        len(brave_res),
        len(ddg_res),
        len(serper_res),
        source_count,
        scrape_count,
        min(scores) if scores else 0.0,
        max(scores) if scores else 0.0,
        response_time,
    )

    response = SearchResponse(
        query=request.query,
        answer=answer,
        results=results,
        response_time=response_time,
        credits_used=2 if request.search_depth == "advanced" else 1,
    )
    await cache.set(cache_key, response.model_dump_json(), ttl=600)
    return response