Import source: axon/routers/search

Import the existing local project snapshot. This commit organizes existing files; it does not represent newly developed functionality.

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