diff --git a/axon/routers/search.py b/axon/routers/search.py new file mode 100644 index 0000000..7b7e1ca --- /dev/null +++ b/axon/routers/search.py @@ -0,0 +1,165 @@ +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