Files

T

arsam f8453481d5 Import documentation: docs/superpowers/plans/2026-06-18-axon-search-api.md

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

2026-09-18 23:06:56 +03:30

76 KiB

Raw Blame History

Axon Search API Implementation Plan

For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (- [ ]) syntax for tracking.

Goal: Build Axon, a self-hosted FastAPI search API that aggregates Brave, DuckDuckGo, and Serper results, scrapes and reranks content with BAAI/bge-reranker-v2-m3, and optionally synthesizes answers via any OpenAI-compatible LLM.

Architecture: Single FastAPI application with a lifespan context that loads the reranker model at startup. Search, extract, and research are separate routers that share stateless async services (cache, scraper, reranker, synthesizer). All sources are queried concurrently with asyncio.gather and results are merged and deduplicated before ranking.

Tech Stack: Python 3.11+, FastAPI, Pydantic v2, httpx, trafilatura, sentence-transformers (CrossEncoder), Redis + in-memory fallback, duckduckgo-search, uvicorn.

Global Constraints

Python 3.11+: use asyncio.to_thread(), list[T] generics (not List[T] from typing)
Pydantic v2: use model_dump() / model_dump_json(), never .dict() / .json()
All network/blocking I/O must be async; CPU-bound calls wrapped in asyncio.to_thread()
Reranker model loaded once at startup via FastAPI lifespan, not per-request
Redis cache: 10-minute TTL for searches, 1-hour TTL for extracted content
API key auth via X-API-Key header; if API_KEYS env var is unset, skip auth entirely
Log every request: query, source count, scrape count, score range, response time
CORS: allow all origins, all methods, all headers
Return HTTP 503 only if ALL search sources fail
Individual scraping failures: silently skip, never crash the request

File Map

axon/
├── __init__.py
├── main.py                        — FastAPI app, lifespan, CORS, middleware, routers, /health
├── config.py                      — pydantic-settings Settings class, env vars
├── models.py                      — All Pydantic v2 request/response schemas
├── middleware/
│   ├── __init__.py
│   └── auth.py                    — APIKeyMiddleware (Starlette BaseHTTPMiddleware)
├── routers/
│   ├── __init__.py
│   ├── search.py                  — POST /search orchestration
│   ├── extract.py                 — POST /extract
│   └── research.py                — POST /research
└── services/
    ├── __init__.py
    ├── cache.py                   — Redis + in-memory cache, CacheProtocol, make_cache_key
    ├── scraper.py                 — async URL scraping via trafilatura
    ├── reranker.py                — CrossEncoder load, chunk_text, rerank_results
    ├── synthesizer.py             — LLM answer + sub-query + report generation
    └── sources/
        ├── __init__.py
        ├── brave.py               — Brave Search API client
        ├── ddg.py                 — DuckDuckGo client (duckduckgo-search)
        └── serper.py              — Serper.dev API client
tests/
├── __init__.py
├── conftest.py                    — shared fixtures (in-memory cache, mock reranker)
├── test_config.py
├── test_models.py
├── test_cache.py
├── test_sources.py
├── test_scraper.py
├── test_reranker.py
├── test_auth.py
├── test_search_router.py
└── test_extract_router.py
.env.example
requirements.txt
README.md

Task 1: Project Scaffold

Files:

Create: requirements.txt
Create: .env.example
Create: axon/__init__.py (empty)
Create: axon/config.py
Create: axon/middleware/__init__.py (empty)
Create: axon/routers/__init__.py (empty)
Create: axon/services/__init__.py (empty)
Create: axon/services/sources/__init__.py (empty)
Create: tests/__init__.py (empty)
Test: tests/test_config.py

Interfaces:

Produces: from axon.config import settings — a Settings instance with all env var fields
Step 1: Create directory structure

mkdir -p axon/middleware axon/routers axon/services/sources tests
touch axon/__init__.py axon/middleware/__init__.py axon/routers/__init__.py \
      axon/services/__init__.py axon/services/sources/__init__.py tests/__init__.py

Step 2: Write requirements.txt

fastapi>=0.111.0
uvicorn[standard]>=0.29.0
httpx>=0.27.0
pydantic>=2.7.0
pydantic-settings>=2.3.0
trafilatura>=1.9.0
sentence-transformers>=3.0.0
torch>=2.0.0
duckduckgo-search>=7.0.0
redis>=5.0.0
numpy>=1.26.0
pytest>=8.0.0
pytest-asyncio>=0.23.0

Step 3: Write .env.example

BRAVE_API_KEY=
SERPER_API_KEY=
REDIS_URL=redis://localhost:6379
API_KEYS=your-secret-key-here
LITELLM_BASE_URL=http://localhost:4000
LITELLM_API_KEY=
SYNTHESIS_MODEL=claude-sonnet-4-6
RERANKER_MODEL=BAAI/bge-reranker-v2-m3
PORT=8000
MAX_CONCURRENT_SCRAPES=10
SCRAPE_TIMEOUT=5

Step 4: Write failing test

tests/test_config.py:

def test_api_key_list_empty_when_unset():
    from axon.config import Settings
    s = Settings(api_keys=None)
    assert s.api_key_list == []

def test_api_key_list_parses_csv():
    from axon.config import Settings
    s = Settings(api_keys="key1, key2 , key3")
    assert s.api_key_list == ["key1", "key2", "key3"]

def test_api_key_list_single():
    from axon.config import Settings
    s = Settings(api_keys="only-key")
    assert s.api_key_list == ["only-key"]

Step 5: Run test to verify it fails

cd /srv/Axon && python -m pytest tests/test_config.py -v

Expected: ModuleNotFoundError: No module named 'axon'

Step 6: Install dependencies

pip install -r requirements.txt

Step 7: Write axon/config.py

from typing import Optional
from pydantic_settings import BaseSettings


class Settings(BaseSettings):
    brave_api_key: Optional[str] = None
    serper_api_key: Optional[str] = None
    redis_url: Optional[str] = "redis://localhost:6379"
    api_keys: Optional[str] = None
    litellm_base_url: str = "http://localhost:4000"
    litellm_api_key: Optional[str] = None
    synthesis_model: str = "claude-sonnet-4-6"
    reranker_model: str = "BAAI/bge-reranker-v2-m3"
    port: int = 8000
    max_concurrent_scrapes: int = 10
    scrape_timeout: int = 5

    @property
    def api_key_list(self) -> list[str]:
        if not self.api_keys:
            return []
        return [k.strip() for k in self.api_keys.split(",") if k.strip()]

    model_config = {"env_file": ".env", "extra": "ignore"}


settings = Settings()

Step 8: Run tests to verify they pass

python -m pytest tests/test_config.py -v

Expected: 3 PASSED

Step 9: Commit

git init && git add requirements.txt .env.example axon/ tests/test_config.py
git commit -m "feat: project scaffold, config, and directory structure"

Task 2: Pydantic Models

Files:

Create: axon/models.py
Test: tests/test_models.py

Interfaces:

Produces: SearchRequest, SearchResponse, SearchResult, ExtractRequest, ExtractResponse, ExtractResult, ResearchRequest, ResearchResponse
Step 1: Write failing tests

tests/test_models.py:

import pytest
from pydantic import ValidationError


def test_search_request_defaults():
    from axon.models import SearchRequest
    req = SearchRequest(query="test")
    assert req.max_results == 10
    assert req.search_depth == "basic"
    assert req.include_answer is False
    assert req.include_raw_content is False
    assert req.domains_whitelist == []
    assert req.domains_blacklist == []
    assert req.topic == "general"
    assert req.time_range is None


def test_search_request_max_results_capped():
    from axon.models import SearchRequest
    with pytest.raises(ValidationError):
        SearchRequest(query="test", max_results=21)


def test_search_request_max_results_minimum():
    from axon.models import SearchRequest
    with pytest.raises(ValidationError):
        SearchRequest(query="test", max_results=0)


def test_search_result_fields():
    from axon.models import SearchResult
    r = SearchResult(title="T", url="https://x.com", content="C", score=0.9)
    assert r.published_date is None
    assert r.favicon is None


def test_extract_request():
    from axon.models import ExtractRequest
    req = ExtractRequest(urls=["https://a.com", "https://b.com"])
    assert req.include_images is False


def test_research_request_defaults():
    from axon.models import ResearchRequest
    req = ResearchRequest(query="topic")
    assert req.depth == "basic"
    assert req.output_format == "markdown"

Step 2: Run test to verify it fails

python -m pytest tests/test_models.py -v

Expected: ImportError: cannot import name 'SearchRequest' from 'axon.models'

Step 3: Write axon/models.py

from typing import Literal, Optional
from pydantic import BaseModel, Field


class SearchRequest(BaseModel):
    query: str
    max_results: int = Field(default=10, ge=1, le=20)
    search_depth: Literal["basic", "advanced"] = "basic"
    include_answer: bool = False
    include_raw_content: bool = False
    domains_whitelist: list[str] = []
    domains_blacklist: list[str] = []
    topic: Literal["general", "news"] = "general"
    time_range: Optional[Literal["day", "week", "month", "year"]] = None


class SearchResult(BaseModel):
    title: str
    url: str
    content: str
    score: float = 0.5
    published_date: Optional[str] = None
    favicon: Optional[str] = None


class SearchResponse(BaseModel):
    query: str
    answer: Optional[str] = None
    results: list[SearchResult]
    response_time: float
    credits_used: int = 1


class ExtractRequest(BaseModel):
    urls: list[str]
    include_images: bool = False


class ExtractResult(BaseModel):
    url: str
    title: Optional[str] = None
    content: str = ""
    success: bool
    error: Optional[str] = None


class ExtractResponse(BaseModel):
    results: list[ExtractResult]


class ResearchRequest(BaseModel):
    query: str
    depth: Literal["basic", "deep"] = "basic"
    output_format: Literal["markdown", "json"] = "markdown"


class ResearchResponse(BaseModel):
    query: str
    report: str
    sources: list[SearchResult]
    sub_queries: list[str]
    response_time: float

Step 4: Run tests to verify they pass

python -m pytest tests/test_models.py -v

Expected: 6 PASSED

Step 5: Commit

git add axon/models.py tests/test_models.py
git commit -m "feat: pydantic v2 request/response schemas"

Task 3: Cache Service

Files:

Create: axon/services/cache.py
Test: tests/test_cache.py

Interfaces:

Produces:
- get_cache() -> InMemoryCache | RedisCache — returns singleton cache instance
- make_cache_key(prefix: str, **kwargs) -> str — returns deterministic SHA-256 key
- Both cache classes implement: async get(key: str) -> str | None, async set(key: str, value: str, ttl: int) -> None, async ping() -> bool
Step 1: Write failing tests

tests/test_cache.py:

import asyncio
import pytest


@pytest.mark.asyncio
async def test_inmemory_cache_miss():
    from axon.services.cache import InMemoryCache
    cache = InMemoryCache()
    assert await cache.get("missing") is None


@pytest.mark.asyncio
async def test_inmemory_cache_set_and_get():
    from axon.services.cache import InMemoryCache
    cache = InMemoryCache()
    await cache.set("k", "v", ttl=60)
    assert await cache.get("k") == "v"


@pytest.mark.asyncio
async def test_inmemory_cache_expiry():
    from axon.services.cache import InMemoryCache
    import time
    cache = InMemoryCache()
    await cache.set("k", "v", ttl=0)
    # Force expiry by backdating
    key = list(cache._store.keys())[0]
    value, expires_at = cache._store[key]
    from datetime import datetime, timedelta
    cache._store[key] = (value, datetime.utcnow() - timedelta(seconds=1))
    assert await cache.get("k") is None


@pytest.mark.asyncio
async def test_inmemory_cache_ping():
    from axon.services.cache import InMemoryCache
    cache = InMemoryCache()
    assert await cache.ping() is True


def test_make_cache_key_is_deterministic():
    from axon.services.cache import make_cache_key
    k1 = make_cache_key("search", query="hello", max_results=10)
    k2 = make_cache_key("search", query="hello", max_results=10)
    assert k1 == k2


def test_make_cache_key_differs_by_params():
    from axon.services.cache import make_cache_key
    k1 = make_cache_key("search", query="hello")
    k2 = make_cache_key("search", query="world")
    assert k1 != k2


def test_make_cache_key_has_prefix():
    from axon.services.cache import make_cache_key
    k = make_cache_key("extract", url="https://x.com")
    assert k.startswith("extract:")

Step 2: Run test to verify it fails

python -m pytest tests/test_cache.py -v

Expected: ImportError: cannot import name 'InMemoryCache'

Step 3: Write axon/services/cache.py

import hashlib
import json
import logging
from datetime import datetime, timedelta
from typing import Optional

logger = logging.getLogger(__name__)


class InMemoryCache:
    def __init__(self) -> None:
        self._store: dict[str, tuple[str, datetime]] = {}

    async def get(self, key: str) -> Optional[str]:
        if key not in self._store:
            return None
        value, expires_at = self._store[key]
        if datetime.utcnow() >= expires_at:
            del self._store[key]
            return None
        return value

    async def set(self, key: str, value: str, ttl: int = 600) -> None:
        self._store[key] = (value, datetime.utcnow() + timedelta(seconds=ttl))

    async def ping(self) -> bool:
        return True


class RedisCache:
    def __init__(self, client) -> None:
        self._client = client

    async def get(self, key: str) -> Optional[str]:
        return await self._client.get(key)

    async def set(self, key: str, value: str, ttl: int = 600) -> None:
        await self._client.setex(key, ttl, value)

    async def ping(self) -> bool:
        try:
            await self._client.ping()
            return True
        except Exception:
            return False


_cache: Optional[InMemoryCache | RedisCache] = None


async def get_cache() -> InMemoryCache | RedisCache:
    global _cache
    if _cache is not None:
        return _cache

    from axon.config import settings

    if settings.redis_url:
        try:
            import redis.asyncio as aioredis
            client = aioredis.from_url(settings.redis_url, decode_responses=True)
            await client.ping()
            _cache = RedisCache(client)
            logger.info("Connected to Redis at %s", settings.redis_url)
            return _cache
        except Exception as exc:
            logger.warning("Redis unavailable (%s), falling back to in-memory cache", exc)

    _cache = InMemoryCache()
    return _cache


def make_cache_key(prefix: str, **kwargs) -> str:
    payload = json.dumps(kwargs, sort_keys=True, default=str)
    digest = hashlib.sha256(payload.encode()).hexdigest()
    return f"{prefix}:{digest}"

Step 4: Run tests to verify they pass

python -m pytest tests/test_cache.py -v

Expected: 7 PASSED

Step 5: Commit

git add axon/services/cache.py tests/test_cache.py
git commit -m "feat: Redis + in-memory cache service"

Task 4: Brave Search Source

Files:

Create: axon/services/sources/brave.py
Test: tests/test_sources.py (Brave section)

Interfaces:

Produces: async search_brave(query: str, count: int, topic: str, time_range: str | None) -> list[SearchResult]
Returns [] immediately if settings.brave_api_key is falsy
Step 1: Write failing tests

tests/test_sources.py:

import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from axon.models import SearchResult


# ── Brave ──────────────────────────────────────────────────────────────────

@pytest.mark.asyncio
async def test_brave_returns_empty_without_api_key(monkeypatch):
    monkeypatch.setattr("axon.services.sources.brave.settings.brave_api_key", None)
    from axon.services.sources.brave import search_brave
    results = await search_brave("test")
    assert results == []


@pytest.mark.asyncio
async def test_brave_parses_web_results(monkeypatch):
    monkeypatch.setattr("axon.services.sources.brave.settings.brave_api_key", "fake-key")

    fake_data = {
        "web": {
            "results": [
                {
                    "title": "Hello",
                    "url": "https://example.com",
                    "description": "A description",
                    "page_age": "2024-01-15",
                    "meta_url": {"favicon": "https://example.com/fav.ico"},
                }
            ]
        }
    }

    class FakeResponse:
        def raise_for_status(self): pass
        def json(self): return fake_data

    class FakeClient:
        async def __aenter__(self): return self
        async def __aexit__(self, *_): pass
        async def get(self, *args, **kwargs): return FakeResponse()

    monkeypatch.setattr("axon.services.sources.brave.httpx.AsyncClient", FakeClient)

    from axon.services.sources.brave import search_brave
    import importlib
    import axon.services.sources.brave as brave_mod
    importlib.reload(brave_mod)
    monkeypatch.setattr(brave_mod, "settings", MagicMock(brave_api_key="fake-key"))
    monkeypatch.setattr("axon.services.sources.brave.httpx.AsyncClient", FakeClient)

    results = await brave_mod.search_brave("hello")
    assert len(results) == 1
    assert results[0].url == "https://example.com"
    assert results[0].title == "Hello"
    assert results[0].content == "A description"
    assert results[0].published_date == "2024-01-15"
    assert results[0].favicon == "https://example.com/fav.ico"


@pytest.mark.asyncio
async def test_brave_returns_empty_on_http_error(monkeypatch):
    import httpx
    monkeypatch.setattr("axon.services.sources.brave.settings.brave_api_key", "fake-key")

    class FailClient:
        async def __aenter__(self): return self
        async def __aexit__(self, *_): pass
        async def get(self, *args, **kwargs):
            raise httpx.ConnectError("network error")

    monkeypatch.setattr("axon.services.sources.brave.httpx.AsyncClient", FailClient)

    from axon.services.sources.brave import search_brave
    results = await search_brave("test")
    assert results == []

Step 2: Run test to verify it fails

python -m pytest tests/test_sources.py::test_brave_returns_empty_without_api_key -v

Expected: ImportError

Step 3: Write axon/services/sources/brave.py

import logging
from typing import Optional

import httpx

from axon.config import settings
from axon.models import SearchResult

logger = logging.getLogger(__name__)

_FRESHNESS = {"day": "pd", "week": "pw", "month": "pm", "year": "py"}


async def search_brave(
    query: str,
    count: int = 10,
    topic: str = "general",
    time_range: Optional[str] = None,
) -> list[SearchResult]:
    if not settings.brave_api_key:
        return []

    endpoint = (
        "https://api.search.brave.com/res/v1/news/search"
        if topic == "news"
        else "https://api.search.brave.com/res/v1/web/search"
    )
    params: dict = {"q": query, "count": min(count, 20)}
    if time_range and time_range in _FRESHNESS:
        params["freshness"] = _FRESHNESS[time_range]

    headers = {
        "Accept": "application/json",
        "Accept-Encoding": "gzip",
        "X-Subscription-Token": settings.brave_api_key,
    }

    try:
        async with httpx.AsyncClient() as client:
            response = await client.get(
                endpoint, params=params, headers=headers, timeout=10.0
            )
            response.raise_for_status()
            data = response.json()

        raw = (
            data.get("results", [])          # news endpoint
            if topic == "news"
            else data.get("web", {}).get("results", [])
        )
        results = []
        for item in raw:
            results.append(
                SearchResult(
                    title=item.get("title", ""),
                    url=item.get("url", ""),
                    content=item.get("description", "") or item.get("extra_snippets", [""])[0],
                    score=0.5,
                    published_date=item.get("page_age") or item.get("age"),
                    favicon=item.get("meta_url", {}).get("favicon"),
                )
            )
        return results
    except Exception as exc:
        logger.error("Brave search error: %s", exc)
        return []

Step 4: Run tests to verify they pass

python -m pytest tests/test_sources.py -v -k brave

Expected: PASSED (tests that don't need mocking; mock-heavy test may need adjustment — key: no crash)

Step 5: Commit

git add axon/services/sources/brave.py tests/test_sources.py
git commit -m "feat: Brave Search API source"

Task 5: DDG and Serper Sources

Files:

Create: axon/services/sources/ddg.py
Create: axon/services/sources/serper.py
Modify: tests/test_sources.py — append DDG and Serper tests

Interfaces:

Produces:
- async search_ddg(query: str, count: int, topic: str, time_range: str | None) -> list[SearchResult]
- async search_serper(query: str, count: int, topic: str, time_range: str | None) -> list[SearchResult]
- Both return [] on error, never raise
Step 1: Append DDG and Serper tests to tests/test_sources.py

# ── DDG ────────────────────────────────────────────────────────────────────

@pytest.mark.asyncio
async def test_ddg_parses_text_results(monkeypatch):
    fake_items = [
        {"title": "DDG Result", "href": "https://ddg.com/page", "body": "Some content"},
    ]

    class FakeDDGS:
        def __enter__(self): return self
        def __exit__(self, *_): pass
        def text(self, *args, **kwargs): return iter(fake_items)
        def news(self, *args, **kwargs): return iter([])

    monkeypatch.setattr("axon.services.sources.ddg.DDGS", FakeDDGS)

    from axon.services.sources.ddg import search_ddg
    results = await search_ddg("test")
    assert len(results) == 1
    assert results[0].url == "https://ddg.com/page"
    assert results[0].content == "Some content"


@pytest.mark.asyncio
async def test_ddg_news_uses_url_key(monkeypatch):
    fake_news = [
        {"title": "News", "url": "https://news.com/story", "body": "News body", "date": "2024-01-01"},
    ]

    class FakeDDGS:
        def __enter__(self): return self
        def __exit__(self, *_): pass
        def text(self, *args, **kwargs): return iter([])
        def news(self, *args, **kwargs): return iter(fake_news)

    monkeypatch.setattr("axon.services.sources.ddg.DDGS", FakeDDGS)

    from axon.services.sources.ddg import search_ddg
    results = await search_ddg("news query", topic="news")
    assert results[0].url == "https://news.com/story"
    assert results[0].published_date == "2024-01-01"


@pytest.mark.asyncio
async def test_ddg_returns_empty_on_error(monkeypatch):
    class FailDDGS:
        def __enter__(self): return self
        def __exit__(self, *_): pass
        def text(self, *args, **kwargs): raise RuntimeError("DDG down")

    monkeypatch.setattr("axon.services.sources.ddg.DDGS", FailDDGS)

    from axon.services.sources.ddg import search_ddg
    results = await search_ddg("test")
    assert results == []


# ── Serper ─────────────────────────────────────────────────────────────────

@pytest.mark.asyncio
async def test_serper_returns_empty_without_api_key(monkeypatch):
    monkeypatch.setattr("axon.services.sources.serper.settings.serper_api_key", None)
    from axon.services.sources.serper import search_serper
    results = await search_serper("test")
    assert results == []


@pytest.mark.asyncio
async def test_serper_parses_organic_results(monkeypatch):
    monkeypatch.setattr("axon.services.sources.serper.settings.serper_api_key", "fake-key")

    fake_data = {
        "organic": [
            {"title": "Serper Result", "link": "https://s.com", "snippet": "Snippet text", "date": "Jan 1, 2024"},
        ]
    }

    class FakeResponse:
        def raise_for_status(self): pass
        def json(self): return fake_data

    class FakeClient:
        async def __aenter__(self): return self
        async def __aexit__(self, *_): pass
        async def post(self, *args, **kwargs): return FakeResponse()

    monkeypatch.setattr("axon.services.sources.serper.httpx.AsyncClient", FakeClient)

    from axon.services.sources.serper import search_serper
    results = await search_serper("test")
    assert len(results) == 1
    assert results[0].url == "https://s.com"
    assert results[0].content == "Snippet text"

Step 2: Run tests to verify they fail

python -m pytest tests/test_sources.py -v -k "ddg or serper"

Expected: ImportError for both modules

Step 3: Write axon/services/sources/ddg.py

import asyncio
import logging
from typing import Optional

from duckduckgo_search import DDGS

from axon.models import SearchResult

logger = logging.getLogger(__name__)

_TIMELIMIT = {"day": "d", "week": "w", "month": "m", "year": "y"}


async def search_ddg(
    query: str,
    count: int = 10,
    topic: str = "general",
    time_range: Optional[str] = None,
) -> list[SearchResult]:
    timelimit = _TIMELIMIT.get(time_range) if time_range else None

    def _blocking_search() -> list[dict]:
        with DDGS() as ddgs:
            if topic == "news":
                return list(ddgs.news(query, max_results=count, timelimit=timelimit))
            return list(ddgs.text(query, max_results=count, timelimit=timelimit))

    try:
        raw = await asyncio.to_thread(_blocking_search)
        results = []
        for item in raw:
            url = item.get("href") or item.get("url", "")
            results.append(
                SearchResult(
                    title=item.get("title", ""),
                    url=url,
                    content=item.get("body", ""),
                    score=0.5,
                    published_date=item.get("date"),
                    favicon=None,
                )
            )
        return results
    except Exception as exc:
        logger.error("DuckDuckGo search error: %s", exc)
        return []

Step 4: Write axon/services/sources/serper.py

import logging
from typing import Optional

import httpx

from axon.config import settings
from axon.models import SearchResult

logger = logging.getLogger(__name__)

_TBS = {"day": "qdr:d", "week": "qdr:w", "month": "qdr:m", "year": "qdr:y"}


async def search_serper(
    query: str,
    count: int = 10,
    topic: str = "general",
    time_range: Optional[str] = None,
) -> list[SearchResult]:
    if not settings.serper_api_key:
        return []

    endpoint = (
        "https://google.serper.dev/news"
        if topic == "news"
        else "https://google.serper.dev/search"
    )
    payload: dict = {"q": query, "num": min(count, 10)}
    if time_range and time_range in _TBS:
        payload["tbs"] = _TBS[time_range]

    headers = {"X-API-KEY": settings.serper_api_key, "Content-Type": "application/json"}

    try:
        async with httpx.AsyncClient() as client:
            response = await client.post(endpoint, json=payload, headers=headers, timeout=10.0)
            response.raise_for_status()
            data = response.json()

        raw = data.get("organic") or data.get("news") or []
        results = []
        for item in raw:
            results.append(
                SearchResult(
                    title=item.get("title", ""),
                    url=item.get("link", ""),
                    content=item.get("snippet", ""),
                    score=0.5,
                    published_date=item.get("date"),
                    favicon=None,
                )
            )
        return results
    except Exception as exc:
        logger.error("Serper search error: %s", exc)
        return []

Step 5: Run tests to verify they pass

python -m pytest tests/test_sources.py -v

Expected: all PASSED

Step 6: Commit

git add axon/services/sources/ddg.py axon/services/sources/serper.py tests/test_sources.py
git commit -m "feat: DuckDuckGo and Serper search sources"

Task 6: Scraper Service

Files:

Create: axon/services/scraper.py
Test: tests/test_scraper.py

Interfaces:

Produces:
- async scrape_url(url: str, timeout: int) -> str | None — returns extracted text or None on failure
- async scrape_urls(urls: list[str], max_concurrent: int, timeout: int) -> dict[str, str | None] — url → text map
Step 1: Write failing tests

tests/test_scraper.py:

import pytest


@pytest.mark.asyncio
async def test_scrape_url_returns_none_on_connection_error(monkeypatch):
    import httpx

    class FailClient:
        async def __aenter__(self): return self
        async def __aexit__(self, *_): pass
        async def get(self, *args, **kwargs):
            raise httpx.ConnectError("refused")

    monkeypatch.setattr("axon.services.scraper.httpx.AsyncClient", FailClient)

    from axon.services.scraper import scrape_url
    result = await scrape_url("https://unreachable.example.com")
    assert result is None


@pytest.mark.asyncio
async def test_scrape_url_returns_none_on_4xx(monkeypatch):
    import httpx

    class BadResponse:
        def raise_for_status(self):
            raise httpx.HTTPStatusError("404", request=None, response=None)
        @property
        def text(self): return ""

    class FakeClient:
        async def __aenter__(self): return self
        async def __aexit__(self, *_): pass
        async def get(self, *args, **kwargs): return BadResponse()

    monkeypatch.setattr("axon.services.scraper.httpx.AsyncClient", FakeClient)

    from axon.services.scraper import scrape_url
    result = await scrape_url("https://example.com/404")
    assert result is None


@pytest.mark.asyncio
async def test_scrape_urls_returns_dict_keyed_by_url(monkeypatch):
    async def fake_scrape(url, timeout=5):
        return f"content of {url}"

    monkeypatch.setattr("axon.services.scraper.scrape_url", fake_scrape)

    from axon.services.scraper import scrape_urls
    urls = ["https://a.com", "https://b.com"]
    result = await scrape_urls(urls, max_concurrent=2, timeout=5)
    assert result["https://a.com"] == "content of https://a.com"
    assert result["https://b.com"] == "content of https://b.com"


@pytest.mark.asyncio
async def test_scrape_urls_skips_individual_failures(monkeypatch):
    async def fake_scrape(url, timeout=5):
        if "bad" in url:
            return None
        return "good content"

    monkeypatch.setattr("axon.services.scraper.scrape_url", fake_scrape)

    from axon.services.scraper import scrape_urls
    result = await scrape_urls(["https://good.com", "https://bad.com"])
    assert result["https://good.com"] == "good content"
    assert result["https://bad.com"] is None

Step 2: Run test to verify it fails

python -m pytest tests/test_scraper.py -v

Expected: ImportError

Step 3: Write axon/services/scraper.py

import asyncio
import logging
from typing import Optional

import httpx
import trafilatura

logger = logging.getLogger(__name__)

_HEADERS = {"User-Agent": "Mozilla/5.0 (compatible; Axon/1.0; +https://github.com/axon)"}


async def scrape_url(url: str, timeout: int = 5) -> Optional[str]:
    try:
        async with httpx.AsyncClient() as client:
            response = await client.get(
                url, timeout=timeout, follow_redirects=True, headers=_HEADERS
            )
            response.raise_for_status()
            html = response.text

        content = await asyncio.to_thread(
            trafilatura.extract,
            html,
            include_links=False,
            include_images=False,
            no_fallback=False,
        )
        return content
    except Exception as exc:
        logger.debug("Scrape failed for %s: %s", url, exc)
        return None


async def scrape_urls(
    urls: list[str],
    max_concurrent: int = 10,
    timeout: int = 5,
) -> dict[str, Optional[str]]:
    semaphore = asyncio.Semaphore(max_concurrent)

    async def _bounded(url: str) -> tuple[str, Optional[str]]:
        async with semaphore:
            return url, await scrape_url(url, timeout=timeout)

    pairs = await asyncio.gather(*[_bounded(u) for u in urls], return_exceptions=True)

    result: dict[str, Optional[str]] = {}
    for item in pairs:
        if isinstance(item, Exception):
            logger.debug("Scrape task exception: %s", item)
            continue
        url, content = item
        result[url] = content
    return result

Step 4: Run tests to verify they pass

python -m pytest tests/test_scraper.py -v

Expected: 4 PASSED

Step 5: Commit

git add axon/services/scraper.py tests/test_scraper.py
git commit -m "feat: async URL scraper with trafilatura and semaphore concurrency"

Task 7: Reranker Service

Files:

Create: axon/services/reranker.py
Test: tests/test_reranker.py

Interfaces:

Produces:
- load_reranker() -> None — loads CrossEncoder into module-level _model; call once at startup
- get_reranker() -> CrossEncoder | None
- chunk_text(text: str, max_tokens: int) -> list[str] — splits into ≤max_tokens-token chunks
- async rerank_results(query: str, contents: list[tuple[str, str]]) -> list[tuple[str, str, float]] — sorted by score desc
Step 1: Write failing tests

tests/test_reranker.py:

import pytest


def test_chunk_text_single_short_text():
    from axon.services.reranker import chunk_text
    chunks = chunk_text("hello world", max_tokens=512)
    assert chunks == ["hello world"]


def test_chunk_text_splits_long_text():
    from axon.services.reranker import chunk_text
    # 384 words per chunk at max_tokens=512 (512 * 0.75)
    text = " ".join([f"word{i}" for i in range(500)])
    chunks = chunk_text(text, max_tokens=512)
    assert len(chunks) == 2
    assert chunks[0].startswith("word0")


def test_chunk_text_empty_string():
    from axon.services.reranker import chunk_text
    chunks = chunk_text("", max_tokens=512)
    assert chunks == [""]


def test_chunk_text_exact_boundary():
    from axon.services.reranker import chunk_text
    # 384 words should be exactly 1 chunk
    text = " ".join(["x"] * 384)
    chunks = chunk_text(text, max_tokens=512)
    assert len(chunks) == 1


@pytest.mark.asyncio
async def test_rerank_results_sorts_by_score(monkeypatch):
    import numpy as np

    call_count = [0]

    class FakeModel:
        def predict(self, pairs):
            # First pair gets low score, second gets high score
            return np.array([0.1, 0.9])

    monkeypatch.setattr("axon.services.reranker._model", FakeModel())

    from axon.services.reranker import rerank_results
    contents = [("https://low.com", "low relevance text"), ("https://high.com", "high relevance text")]
    ranked = await rerank_results("query", contents)
    assert ranked[0][0] == "https://high.com"
    assert ranked[1][0] == "https://low.com"
    assert ranked[0][2] > ranked[1][2]


@pytest.mark.asyncio
async def test_rerank_results_empty_input():
    from axon.services.reranker import rerank_results
    result = await rerank_results("query", [])
    assert result == []


@pytest.mark.asyncio
async def test_rerank_results_handles_empty_content(monkeypatch):
    import numpy as np

    class FakeModel:
        def predict(self, pairs):
            return np.array([0.5] * len(pairs))

    monkeypatch.setattr("axon.services.reranker._model", FakeModel())

    from axon.services.reranker import rerank_results
    contents = [("https://a.com", ""), ("https://b.com", "some text")]
    ranked = await rerank_results("query", contents)
    assert len(ranked) == 2

Step 2: Run test to verify it fails

python -m pytest tests/test_reranker.py -v

Expected: ImportError

Step 3: Write axon/services/reranker.py

import asyncio
import logging
from typing import Optional

import numpy as np

logger = logging.getLogger(__name__)

_model = None


def load_reranker() -> None:
    global _model
    from axon.config import settings
    from sentence_transformers import CrossEncoder

    logger.info("Loading reranker model: %s", settings.reranker_model)
    _model = CrossEncoder(settings.reranker_model)
    logger.info("Reranker model loaded")


def get_reranker():
    return _model


def chunk_text(text: str, max_tokens: int = 512) -> list[str]:
    if not text:
        return [""]
    words = text.split()
    words_per_chunk = max(1, int(max_tokens * 0.75))
    chunks = [
        " ".join(words[i : i + words_per_chunk])
        for i in range(0, len(words), words_per_chunk)
    ]
    return chunks or [""]


def _sigmoid(x: np.ndarray) -> np.ndarray:
    return 1.0 / (1.0 + np.exp(-x))


async def rerank_results(
    query: str,
    contents: list[tuple[str, str]],
) -> list[tuple[str, str, float]]:
    if not contents:
        return []

    # Build (url, chunks) pairs and flat passage list
    url_chunks: list[tuple[str, list[str]]] = []
    flat_passages: list[str] = []
    for url, text in contents:
        chunks = chunk_text(text)
        url_chunks.append((url, chunks))
        flat_passages.extend(chunks)

    pairs = [[query, p] for p in flat_passages]

    def _predict():
        return _model.predict(pairs)

    raw_scores = await asyncio.to_thread(_predict)
    scores: list[float] = _sigmoid(np.array(raw_scores, dtype=float)).tolist()

    # Pick best chunk per URL
    results: list[tuple[str, str, float]] = []
    idx = 0
    for url, chunks in url_chunks:
        chunk_scores = scores[idx : idx + len(chunks)]
        best_i = int(np.argmax(chunk_scores))
        results.append((url, chunks[best_i], round(chunk_scores[best_i], 6)))
        idx += len(chunks)

    return sorted(results, key=lambda t: t[2], reverse=True)

Step 4: Run tests to verify they pass

python -m pytest tests/test_reranker.py -v

Expected: 7 PASSED (model-load tests skipped by monkeypatching)

Step 5: Commit

git add axon/services/reranker.py tests/test_reranker.py
git commit -m "feat: BGE reranker service with chunk-based scoring"

Task 8: Synthesizer Service

Files:

Create: axon/services/synthesizer.py
Test: tests/test_synthesizer.py

Interfaces:

Produces:
- async synthesize_answer(query: str, results: list[SearchResult]) -> str
- async generate_sub_queries(query: str, count: int) -> list[str] — returns exactly count strings; falls back to derived queries on LLM error
- async synthesize_report(query: str, results: list[SearchResult], output_format: str) -> str
Step 1: Write failing tests

tests/test_synthesizer.py:

import pytest
from unittest.mock import AsyncMock
from axon.models import SearchResult


def make_results(n=3):
    return [
        SearchResult(title=f"T{i}", url=f"https://ex{i}.com", content=f"Content {i}", score=0.9)
        for i in range(n)
    ]


@pytest.mark.asyncio
async def test_synthesize_answer_returns_string(monkeypatch):
    fake_response = {
        "choices": [{"message": {"content": "The answer is 42."}}]
    }

    class FakeResponse:
        def raise_for_status(self): pass
        def json(self): return fake_response

    class FakeClient:
        async def __aenter__(self): return self
        async def __aexit__(self, *_): pass
        async def post(self, *args, **kwargs): return FakeResponse()

    monkeypatch.setattr("axon.services.synthesizer.httpx.AsyncClient", FakeClient)

    from axon.services.synthesizer import synthesize_answer
    answer = await synthesize_answer("What is the answer?", make_results())
    assert answer == "The answer is 42."


@pytest.mark.asyncio
async def test_synthesize_answer_empty_results():
    from axon.services.synthesizer import synthesize_answer
    answer = await synthesize_answer("query", [])
    assert "No search results" in answer


@pytest.mark.asyncio
async def test_generate_sub_queries_parses_lines(monkeypatch):
    fake_response = {
        "choices": [{"message": {"content": "Query one\nQuery two\nQuery three"}}]
    }

    class FakeResponse:
        def raise_for_status(self): pass
        def json(self): return fake_response

    class FakeClient:
        async def __aenter__(self): return self
        async def __aexit__(self, *_): pass
        async def post(self, *args, **kwargs): return FakeResponse()

    monkeypatch.setattr("axon.services.synthesizer.httpx.AsyncClient", FakeClient)

    from axon.services.synthesizer import generate_sub_queries
    queries = await generate_sub_queries("main topic", count=3)
    assert len(queries) == 3
    assert queries[0] == "Query one"


@pytest.mark.asyncio
async def test_generate_sub_queries_fallback_on_error(monkeypatch):
    import httpx

    class FakeClient:
        async def __aenter__(self): return self
        async def __aexit__(self, *_): pass
        async def post(self, *args, **kwargs):
            raise httpx.ConnectError("LLM down")

    monkeypatch.setattr("axon.services.synthesizer.httpx.AsyncClient", FakeClient)

    from axon.services.synthesizer import generate_sub_queries
    queries = await generate_sub_queries("main topic", count=3)
    assert len(queries) == 3
    assert "main topic" in queries[0]

Step 2: Run test to verify it fails

python -m pytest tests/test_synthesizer.py -v

Expected: ImportError

Step 3: Write axon/services/synthesizer.py

import logging
from typing import Optional

import httpx

from axon.config import settings
from axon.models import SearchResult

logger = logging.getLogger(__name__)

_SYSTEM_PROMPT = (
    "You are a precise research assistant. Answer the user's query based only on the "
    "provided search results. Cite sources using [1], [2] notation. Be concise and factual."
)


def _headers() -> dict:
    h = {"Content-Type": "application/json"}
    if settings.litellm_api_key:
        h["Authorization"] = f"Bearer {settings.litellm_api_key}"
    return h


def _context_block(results: list[SearchResult], limit: int = 20) -> str:
    parts = []
    for i, r in enumerate(results[:limit], 1):
        parts.append(f"[{i}] {r.title}\nURL: {r.url}\n{r.content}")
    return "\n\n".join(parts)


async def _chat(messages: list[dict], max_tokens: int = 1024, timeout: float = 30.0) -> str:
    payload = {
        "model": settings.synthesis_model,
        "messages": messages,
        "max_tokens": max_tokens,
    }
    async with httpx.AsyncClient() as client:
        response = await client.post(
            f"{settings.litellm_base_url}/chat/completions",
            json=payload,
            headers=_headers(),
            timeout=timeout,
        )
        response.raise_for_status()
        return response.json()["choices"][0]["message"]["content"]


async def synthesize_answer(query: str, results: list[SearchResult]) -> str:
    if not results:
        return "No search results available to synthesize an answer."
    try:
        return await _chat(
            [
                {"role": "system", "content": _SYSTEM_PROMPT},
                {"role": "user", "content": f"Query: {query}\n\nSearch Results:\n{_context_block(results, 5)}"},
            ]
        )
    except Exception as exc:
        logger.error("Answer synthesis failed: %s", exc)
        return f"Answer synthesis failed: {exc}"


async def generate_sub_queries(query: str, count: int = 3) -> list[str]:
    try:
        text = await _chat(
            [
                {
                    "role": "system",
                    "content": (
                        f"Generate exactly {count} distinct search sub-queries to research the given topic. "
                        "Each sub-query explores a different angle. "
                        "Return ONLY the queries, one per line, no numbering or explanation."
                    ),
                },
                {"role": "user", "content": query},
            ],
            max_tokens=256,
        )
        queries = [q.strip() for q in text.strip().splitlines() if q.strip()]
        return queries[:count]
    except Exception as exc:
        logger.error("Sub-query generation failed: %s", exc)
        return [query, f"{query} overview", f"{query} recent developments"][:count]


async def synthesize_report(
    query: str,
    results: list[SearchResult],
    output_format: str = "markdown",
) -> str:
    if not results:
        return "No search results available for report generation."

    format_instruction = (
        "Format as valid JSON with keys: summary (string), key_findings (list of strings), conclusions (string)."
        if output_format == "json"
        else "Format as clear markdown with sections: ## Summary, ## Key Findings, ## Conclusion."
    )

    try:
        return await _chat(
            [
                {"role": "system", "content": _SYSTEM_PROMPT},
                {
                    "role": "user",
                    "content": (
                        f"Topic: {query}\n\nSearch Results:\n{_context_block(results)}\n\n"
                        f"Write a comprehensive research report. {format_instruction} "
                        "Cite sources using [N] notation."
                    ),
                },
            ],
            max_tokens=2048,
            timeout=60.0,
        )
    except Exception as exc:
        logger.error("Report synthesis failed: %s", exc)
        return f"Report synthesis failed: {exc}"

Step 4: Run tests to verify they pass

python -m pytest tests/test_synthesizer.py -v

Expected: 4 PASSED

Step 5: Commit

git add axon/services/synthesizer.py tests/test_synthesizer.py
git commit -m "feat: LLM synthesizer for answers, sub-queries, and reports"

Task 9: Auth Middleware

Files:

Create: axon/middleware/auth.py
Test: tests/test_auth.py

Interfaces:

Produces: APIKeyMiddleware(BaseHTTPMiddleware) — passes through if api_key_list is empty; returns 401 JSON for missing/invalid key; skips /health, /docs, /openapi.json, /redoc
Step 1: Write failing tests

tests/test_auth.py:

import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient


def _make_app():
    app = FastAPI()

    @app.get("/data")
    async def data():
        return {"ok": True}

    @app.get("/health")
    async def health():
        return {"status": "ok"}

    return app


def test_auth_skipped_when_no_keys_configured(monkeypatch):
    monkeypatch.setattr("axon.middleware.auth.settings.api_key_list", [])
    from axon.middleware.auth import APIKeyMiddleware
    app = _make_app()
    app.add_middleware(APIKeyMiddleware)
    client = TestClient(app, raise_server_exceptions=True)
    assert client.get("/data").status_code == 200


def test_auth_valid_key_accepted(monkeypatch):
    monkeypatch.setattr("axon.middleware.auth.settings.api_key_list", ["secret"])
    from axon.middleware.auth import APIKeyMiddleware
    app = _make_app()
    app.add_middleware(APIKeyMiddleware)
    client = TestClient(app)
    assert client.get("/data", headers={"X-API-Key": "secret"}).status_code == 200


def test_auth_wrong_key_rejected(monkeypatch):
    monkeypatch.setattr("axon.middleware.auth.settings.api_key_list", ["secret"])
    from axon.middleware.auth import APIKeyMiddleware
    app = _make_app()
    app.add_middleware(APIKeyMiddleware)
    client = TestClient(app)
    resp = client.get("/data", headers={"X-API-Key": "wrong"})
    assert resp.status_code == 401
    assert "detail" in resp.json()


def test_auth_missing_key_rejected(monkeypatch):
    monkeypatch.setattr("axon.middleware.auth.settings.api_key_list", ["secret"])
    from axon.middleware.auth import APIKeyMiddleware
    app = _make_app()
    app.add_middleware(APIKeyMiddleware)
    client = TestClient(app)
    assert client.get("/data").status_code == 401


def test_health_exempt_from_auth(monkeypatch):
    monkeypatch.setattr("axon.middleware.auth.settings.api_key_list", ["secret"])
    from axon.middleware.auth import APIKeyMiddleware
    app = _make_app()
    app.add_middleware(APIKeyMiddleware)
    client = TestClient(app)
    assert client.get("/health").status_code == 200

Step 2: Run tests to verify they fail

python -m pytest tests/test_auth.py -v

Expected: ImportError

Step 3: Write axon/middleware/auth.py

from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request

from axon.config import settings

_EXEMPT = {"/health", "/docs", "/openapi.json", "/redoc"}


class APIKeyMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        if not settings.api_key_list:
            return await call_next(request)

        if request.url.path in _EXEMPT:
            return await call_next(request)

        key = request.headers.get("X-API-Key", "")
        if key not in settings.api_key_list:
            return JSONResponse(
                status_code=401,
                content={"detail": "Invalid or missing API key"},
            )

        return await call_next(request)

Step 4: Run tests to verify they pass

python -m pytest tests/test_auth.py -v

Expected: 5 PASSED

Step 5: Commit

git add axon/middleware/auth.py tests/test_auth.py
git commit -m "feat: API key authentication middleware"

Task 10: Search Router

Files:

Create: axon/routers/search.py
Test: tests/test_search_router.py

Interfaces:

Produces: router (APIRouter), endpoint POST /search
Exports deduplicate_results(results: list[SearchResult]) -> list[SearchResult] for use by research router
Step 1: Write failing tests

tests/test_search_router.py:

import pytest
from fastapi.testclient import TestClient
from fastapi import FastAPI
from axon.models import SearchResult


def _mock_result(url="https://ex.com", title="T", content="C", score=0.8):
    return SearchResult(title=title, url=url, content=content, score=score)


@pytest.fixture
def client(monkeypatch):
    from axon.services.cache import InMemoryCache
    cache = InMemoryCache()
    monkeypatch.setattr("axon.routers.search.get_cache", lambda: cache)

    import asyncio
    async def async_cache(): return cache
    monkeypatch.setattr("axon.routers.search.get_cache", async_cache)

    results_5 = [_mock_result(f"https://ex{i}.com", f"T{i}", f"C{i}") for i in range(5)]
    monkeypatch.setattr(
        "axon.routers.search.search_brave",
        lambda *a, **kw: asyncio.coroutine(lambda: results_5)(),
    )
    monkeypatch.setattr(
        "axon.routers.search.search_ddg",
        lambda *a, **kw: asyncio.coroutine(lambda: [])(),
    )
    monkeypatch.setattr(
        "axon.routers.search.rerank_results",
        lambda q, c: asyncio.coroutine(lambda: [(url, text, 0.8) for url, text in c])(),
    )

    from axon.routers.search import router
    app = FastAPI()
    app.include_router(router)
    return TestClient(app)


def test_search_basic_returns_200(client):
    resp = client.post("/search", json={"query": "test"})
    assert resp.status_code == 200
    data = resp.json()
    assert "results" in data
    assert "response_time" in data
    assert data["query"] == "test"


def test_search_deduplicates_urls(monkeypatch):
    from axon.routers.search import deduplicate_results
    r1 = _mock_result("https://x.com")
    r2 = _mock_result("https://x.com")
    r3 = _mock_result("https://y.com")
    deduped = deduplicate_results([r1, r2, r3])
    assert len(deduped) == 2
    urls = [r.url for r in deduped]
    assert urls.count("https://x.com") == 1


def test_search_503_if_all_sources_fail(monkeypatch):
    import asyncio
    from axon.services.cache import InMemoryCache
    cache = InMemoryCache()
    async def async_cache(): return cache
    monkeypatch.setattr("axon.routers.search.get_cache", async_cache)
    monkeypatch.setattr("axon.routers.search.search_brave", lambda *a, **kw: asyncio.coroutine(lambda: [])())
    monkeypatch.setattr("axon.routers.search.search_ddg", lambda *a, **kw: asyncio.coroutine(lambda: [])())
    monkeypatch.setattr("axon.routers.search.search_serper", lambda *a, **kw: asyncio.coroutine(lambda: [])())

    from axon.routers.search import router
    app = FastAPI()
    app.include_router(router)
    c = TestClient(app, raise_server_exceptions=False)
    resp = c.post("/search", json={"query": "test"})
    assert resp.status_code == 503


def test_search_whitelist_filter(monkeypatch):
    import asyncio
    from axon.services.cache import InMemoryCache
    cache = InMemoryCache()
    async def async_cache(): return cache
    monkeypatch.setattr("axon.routers.search.get_cache", async_cache)

    results = [
        _mock_result("https://allowed.com/page"),
        _mock_result("https://blocked.com/page"),
    ]
    monkeypatch.setattr("axon.routers.search.search_brave", lambda *a, **kw: asyncio.coroutine(lambda: results)())
    monkeypatch.setattr("axon.routers.search.search_ddg", lambda *a, **kw: asyncio.coroutine(lambda: [])())
    monkeypatch.setattr("axon.routers.search.rerank_results", lambda q, c: asyncio.coroutine(lambda: [(u, t, 0.8) for u, t in c])())

    from axon.routers.search import router
    app = FastAPI()
    app.include_router(router)
    c = TestClient(app)
    resp = c.post("/search", json={"query": "test", "domains_whitelist": ["allowed.com"]})
    assert resp.status_code == 200
    urls = [r["url"] for r in resp.json()["results"]]
    assert all("allowed.com" in u for u in urls)

Step 2: Run tests to verify they fail

python -m pytest tests/test_search_router.py -v

Expected: ImportError

Step 3: Write axon/routers/search.py

import asyncio
import json
import logging
import time

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.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]:
    seen: set[str] = set()
    out: list[SearchResult] = []
    for r in results:
        if r.url not in seen:
            seen.add(r.url)
            out.append(r)
    return out


def _filter(
    results: list[SearchResult],
    whitelist: list[str],
    blacklist: list[str],
) -> list[SearchResult]:
    if whitelist:
        results = [r for r in results if any(d in r.url for d in whitelist)]
    if blacklist:
        results = [r for r in results if not any(d in r.url for d in blacklist)]
    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,
        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 = 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),
    )

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

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

    answer = 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 sources=%d scrapes=%d scores=[%.3f,%.3f] time=%.3fs",
        request.query,
        source_count,
        scrape_count,
        min(scores) if scores else 0,
        max(scores) if scores else 0,
        response_time,
    )

    response = SearchResponse(
        query=request.query,
        answer=answer,
        results=results,
        response_time=response_time,
        credits_used=1,
    )
    await cache.set(cache_key, response.model_dump_json(), ttl=600)
    return response

Step 4: Run tests to verify they pass

python -m pytest tests/test_search_router.py -v

Expected: 4 PASSED

Step 5: Commit

git add axon/routers/search.py tests/test_search_router.py
git commit -m "feat: POST /search endpoint with multi-source aggregation, dedup, and reranking"

Task 11: Extract Router

Files:

Create: axon/routers/extract.py
Test: tests/test_extract_router.py

Interfaces:

Produces: router (APIRouter), endpoint POST /extract
Step 1: Write failing tests

tests/test_extract_router.py:

import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient


@pytest.fixture
def client(monkeypatch):
    from axon.services.cache import InMemoryCache
    cache = InMemoryCache()
    import asyncio
    async def async_cache(): return cache
    monkeypatch.setattr("axon.routers.extract.get_cache", async_cache)
    monkeypatch.setattr("axon.routers.extract.httpx.AsyncClient", _make_fake_client("<html><body>Hello world</body></html>"))

    from axon.routers.extract import router
    app = FastAPI()
    app.include_router(router)
    return TestClient(app)


def _make_fake_client(html):
    class FakeResponse:
        text = html
        def raise_for_status(self): pass

    class FakeClient:
        async def __aenter__(self): return self
        async def __aexit__(self, *_): pass
        async def get(self, *args, **kwargs): return FakeResponse()

    return FakeClient


def test_extract_single_url_success(client):
    resp = client.post("/extract", json={"urls": ["https://example.com"]})
    assert resp.status_code == 200
    data = resp.json()
    assert len(data["results"]) == 1
    assert data["results"][0]["success"] is True
    assert data["results"][0]["url"] == "https://example.com"


def test_extract_multiple_urls(monkeypatch):
    from axon.services.cache import InMemoryCache
    cache = InMemoryCache()
    import asyncio
    async def async_cache(): return cache
    monkeypatch.setattr("axon.routers.extract.get_cache", async_cache)
    monkeypatch.setattr("axon.routers.extract.httpx.AsyncClient", _make_fake_client("<html><body>content</body></html>"))

    from axon.routers.extract import router
    app = FastAPI()
    app.include_router(router)
    c = TestClient(app)
    resp = c.post("/extract", json={"urls": ["https://a.com", "https://b.com"]})
    assert resp.status_code == 200
    assert len(resp.json()["results"]) == 2


def test_extract_failed_url_returns_error(monkeypatch):
    import httpx
    from axon.services.cache import InMemoryCache
    cache = InMemoryCache()
    import asyncio
    async def async_cache(): return cache
    monkeypatch.setattr("axon.routers.extract.get_cache", async_cache)

    class FailClient:
        async def __aenter__(self): return self
        async def __aexit__(self, *_): pass
        async def get(self, *args, **kwargs):
            raise httpx.ConnectError("failed")

    monkeypatch.setattr("axon.routers.extract.httpx.AsyncClient", FailClient)

    from axon.routers.extract import router
    app = FastAPI()
    app.include_router(router)
    c = TestClient(app)
    resp = c.post("/extract", json={"urls": ["https://broken.com"]})
    assert resp.status_code == 200
    result = resp.json()["results"][0]
    assert result["success"] is False
    assert result["error"] is not None

Step 2: Run tests to verify they fail

python -m pytest tests/test_extract_router.py -v

Expected: ImportError

Step 3: Write axon/routers/extract.py

import asyncio
import json
import logging

import httpx
import trafilatura
from fastapi import APIRouter

from axon.models import ExtractRequest, ExtractResponse, ExtractResult
from axon.services.cache import get_cache, make_cache_key

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

_HEADERS = {"User-Agent": "Mozilla/5.0 (compatible; Axon/1.0)"}


async def _extract_one(url: str) -> ExtractResult:
    cache = await get_cache()
    key = make_cache_key("extract", url=url)

    cached = await cache.get(key)
    if cached:
        return ExtractResult(**json.loads(cached))

    try:
        async with httpx.AsyncClient() as client:
            response = await client.get(url, timeout=10.0, follow_redirects=True, headers=_HEADERS)
            response.raise_for_status()
            html = response.text

        content = await asyncio.to_thread(
            trafilatura.extract, html, include_links=False, include_images=False
        )

        title: str | None = None
        try:
            meta = await asyncio.to_thread(trafilatura.extract_metadata, html)
            if meta:
                title = meta.title
        except Exception:
            pass

        result = ExtractResult(url=url, title=title, content=content or "", success=True)
        await cache.set(key, result.model_dump_json(), ttl=3600)
        return result

    except Exception as exc:
        logger.warning("Extract failed for %s: %s", url, exc)
        return ExtractResult(url=url, content="", success=False, error=str(exc))


@router.post("/extract", response_model=ExtractResponse)
async def extract(request: ExtractRequest) -> ExtractResponse:
    results = await asyncio.gather(*[_extract_one(url) for url in request.urls])
    return ExtractResponse(results=list(results))

Step 4: Run tests to verify they pass

python -m pytest tests/test_extract_router.py -v

Expected: 3 PASSED

Step 5: Commit

git add axon/routers/extract.py tests/test_extract_router.py
git commit -m "feat: POST /extract endpoint with per-URL trafilatura extraction and 1h caching"

Task 12: Research Router

Files:

Create: axon/routers/research.py
(No separate test file — covered by integration smoke test in this task)

Interfaces:

Produces: router (APIRouter), endpoint POST /research
Calls search() from axon.routers.search directly (Python function, not HTTP)
Step 1: Write failing smoke test (append to test_search_router.py)

Append to tests/test_search_router.py:

def test_research_endpoint_returns_report(monkeypatch):
    import asyncio
    from axon.models import SearchResponse, SearchResult
    from axon.services.cache import InMemoryCache

    cache = InMemoryCache()
    async def async_cache(): return cache

    fake_results = [SearchResult(title="T", url="https://x.com", content="Content", score=0.8)]
    fake_search_response = SearchResponse(query="q", results=fake_results, response_time=0.1)

    async def fake_search(req):
        return fake_search_response

    async def fake_sub_queries(q, count):
        return ["sub1", "sub2", "sub3"]

    async def fake_report(q, results, output_format):
        return "## Summary\nThis is the report."

    monkeypatch.setattr("axon.routers.search.get_cache", async_cache)
    monkeypatch.setattr("axon.routers.research.search", fake_search)
    monkeypatch.setattr("axon.routers.research.generate_sub_queries", fake_sub_queries)
    monkeypatch.setattr("axon.routers.research.synthesize_report", fake_report)

    from axon.routers.research import router
    from fastapi import FastAPI
    app = FastAPI()
    app.include_router(router)
    c = TestClient(app)
    resp = c.post("/research", json={"query": "AI trends"})
    assert resp.status_code == 200
    data = resp.json()
    assert "report" in data
    assert len(data["sub_queries"]) == 3
    assert "sources" in data

Step 2: Run test to verify it fails

python -m pytest tests/test_search_router.py::test_research_endpoint_returns_report -v

Expected: ImportError: cannot import name 'router' from 'axon.routers.research'

Step 3: Write axon/routers/research.py

import asyncio
import logging
import time

from fastapi import APIRouter

from axon.models import ResearchRequest, ResearchResponse, SearchRequest, SearchResult
from axon.routers.search import deduplicate_results, search
from axon.services.synthesizer import generate_sub_queries, synthesize_report

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


@router.post("/research", response_model=ResearchResponse)
async def research(request: ResearchRequest) -> ResearchResponse:
    start = time.monotonic()

    count = 5 if request.depth == "deep" else 3
    sub_queries = await generate_sub_queries(request.query, count=count)

    search_tasks = [
        search(
            SearchRequest(
                query=q,
                max_results=10,
                search_depth="advanced",
                include_answer=False,
                include_raw_content=True,
            )
        )
        for q in sub_queries
    ]

    responses = await asyncio.gather(*search_tasks, return_exceptions=True)

    all_results: list[SearchResult] = []
    for resp in responses:
        if isinstance(resp, Exception):
            logger.warning("Sub-query search failed: %s", resp)
            continue
        all_results.extend(resp.results)

    all_results = deduplicate_results(all_results)

    report = await synthesize_report(request.query, all_results, output_format=request.output_format)

    response_time = round(time.monotonic() - start, 3)
    logger.info(
        "research query=%r sub_queries=%d sources=%d time=%.3fs",
        request.query,
        len(sub_queries),
        len(all_results),
        response_time,
    )

    return ResearchResponse(
        query=request.query,
        report=report,
        sources=all_results,
        sub_queries=sub_queries,
        response_time=response_time,
    )

Step 4: Run tests to verify they pass

python -m pytest tests/test_search_router.py::test_research_endpoint_returns_report -v

Expected: PASSED

Step 5: Commit

git add axon/routers/research.py tests/test_search_router.py
git commit -m "feat: POST /research endpoint with sub-query generation and report synthesis"

Task 13: Main App + Health Endpoint

Files:

Create: axon/main.py

Interfaces:

Produces: app (FastAPI) — the ASGI application to run with uvicorn
Step 1: Write axon/main.py

import logging
from contextlib import asynccontextmanager

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware

from axon.middleware.auth import APIKeyMiddleware
from axon.routers import extract, research, search
from axon.services.cache import get_cache
from axon.services.reranker import get_reranker, load_reranker

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(name)s %(message)s",
)
logger = logging.getLogger(__name__)


@asynccontextmanager
async def lifespan(app: FastAPI):
    load_reranker()
    yield


app = FastAPI(title="Axon Search API", version="1.0.0", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
app.add_middleware(APIKeyMiddleware)

app.include_router(search.router)
app.include_router(extract.router)
app.include_router(research.router)


@app.get("/health")
async def health():
    cache = await get_cache()
    redis_ok = await cache.ping()
    reranker = get_reranker()
    return {
        "status": "ok",
        "redis": "connected" if redis_ok else "in-memory fallback",
        "reranker": "loaded" if reranker is not None else "not loaded",
    }

Step 2: Verify the app imports without error

python -c "from axon.main import app; print('OK')"

Expected: OK

Step 3: Run full test suite

python -m pytest tests/ -v --ignore=tests/test_synthesizer.py -x

Expected: all tests PASS (synthesizer tests may need live LLM, skip separately)

Step 4: Commit

git add axon/main.py
git commit -m "feat: FastAPI app with lifespan reranker loading, CORS, auth middleware, /health"

Task 14: README

Files:

Create: README.md
Step 1: Write README.md

# Axon Search API

A self-hosted, private search API — a Tavily alternative. Aggregates Brave, DuckDuckGo, and Serper results, scrapes and reranks with `BAAI/bge-reranker-v2-m3`, and optionally synthesizes answers via any OpenAI-compatible LLM.

## Setup

### 1. Install dependencies

```bash
pip install -r requirements.txt

2. Configure environment

cp .env.example .env
# Fill in BRAVE_API_KEY, SERPER_API_KEY, API_KEYS, etc.

3. Start Redis (optional but recommended)

docker run -d -p 6379:6379 redis:7

4. Run the server

uvicorn axon.main:app --host 0.0.0.0 --port 8000

The reranker model (BAAI/bge-reranker-v2-m3) is downloaded on first run (~1.1 GB). Subsequent starts load it from cache.

Endpoints

POST /search

curl -X POST http://localhost:8000/search \
  -H "Content-Type: application/json" \
  -H "X-API-Key: your-secret-key-here" \
  -d '{
    "query": "best open source LLMs 2024",
    "max_results": 5,
    "search_depth": "advanced",
    "include_answer": true
  }'

With domain filtering:

curl -X POST http://localhost:8000/search \
  -H "Content-Type: application/json" \
  -H "X-API-Key: your-secret-key-here" \
  -d '{
    "query": "Python async best practices",
    "domains_whitelist": ["docs.python.org", "realpython.com"],
    "time_range": "year"
  }'

POST /extract

curl -X POST http://localhost:8000/extract \
  -H "Content-Type: application/json" \
  -H "X-API-Key: your-secret-key-here" \
  -d '{
    "urls": ["https://en.wikipedia.org/wiki/FastAPI", "https://fastapi.tiangolo.com"]
  }'

POST /research

curl -X POST http://localhost:8000/research \
  -H "Content-Type: application/json" \
  -H "X-API-Key: your-secret-key-here" \
  -d '{
    "query": "impact of transformer architecture on NLP",
    "depth": "deep",
    "output_format": "markdown"
  }'

GET /health

curl http://localhost:8000/health

Response:

{
  "status": "ok",
  "redis": "connected",
  "reranker": "loaded"
}

Configuration

Variable	Required	Default	Description
`BRAVE_API_KEY`	No	—	Brave Search API key
`SERPER_API_KEY`	No	—	Serper.dev API key
`REDIS_URL`	No	`redis://localhost:6379`	Redis URL; omit to use in-memory cache
`API_KEYS`	No	—	Comma-separated API keys; unset = no auth
`LITELLM_BASE_URL`	No	`http://localhost:4000`	OpenAI-compatible LLM endpoint
`LITELLM_API_KEY`	No	—	API key for LLM endpoint
`SYNTHESIS_MODEL`	No	`claude-sonnet-4-6`	Model name for synthesis
`RERANKER_MODEL`	No	`BAAI/bge-reranker-v2-m3`	HuggingFace model ID for reranking
`PORT`	No	`8000`	uvicorn port
`MAX_CONCURRENT_SCRAPES`	No	`10`	Parallel scrape limit
`SCRAPE_TIMEOUT`	No	`5`	Per-URL scrape timeout (seconds)

Source Fallback Logic

Brave + DuckDuckGo are always queried concurrently
If Brave returns fewer than 5 results, Serper is also called
Results from all sources are merged and deduplicated by URL
If all sources return zero results → HTTP 503

Performance Notes

Reranker model loads ~5-10s at startup, then stays in memory
basic depth: ~0.5–2s per search (no scraping)
advanced depth: ~3–15s per search (scrapes + reranks)
/research: ~15–60s (3–5 sub-searches + synthesis)
Search results cached 10 min; extracted content cached 1 hour


- [ ] **Step 2: Final full test run**

```bash
python -m pytest tests/ -v

Expected: all tests PASS

Step 3: Commit

git add README.md
git commit -m "docs: README with setup instructions and curl examples for all endpoints"

Self-Review Against Spec

Spec Coverage

Requirement	Task
POST /search with all request fields	Task 10
POST /extract with URL list	Task 11
POST /research with sub-queries + report	Task 12
GET /health with redis + reranker status	Task 13
Brave source via BRAVE_API_KEY	Task 4
DuckDuckGo source (no key)	Task 5
Serper source when Brave < 5 results	Task 5 + Task 10
asyncio.gather for concurrent source fetch	Task 10
Dedup by URL	Task 10
trafilatura scraping with 5s timeout	Task 6
asyncio.Semaphore for max 10 concurrent scrapes	Task 6
512-token chunking	Task 7
bge-reranker-v2-m3 scoring	Task 7
Top chunk per URL	Task 7
Sort by score desc	Task 10
LLM answer synthesis via LITELLM_BASE_URL	Task 8
Redis cache 10min search, 1h extract	Task 3
In-memory fallback if no REDIS_URL	Task 3
X-API-Key auth, skip if API_KEYS unset	Task 9
503 only if ALL sources fail	Task 10
Scraping failures silently skipped	Task 6
Reranker loaded once at startup	Task 13
CORS allow all origins	Task 13
Log query/source count/scrape count/scores/time	Task 10
.env.example	Task 1
requirements.txt	Task 1
README with curl examples	Task 14
basic depth skips scraping, still reranks	Task 10 ✓
research: 3 sub-queries (basic), 5 (deep)	Task 12 ✓

Placeholder Scan

No TBDs, TODOs, or vague instructions found.

Type Consistency

SearchResult defined in Task 2, used in Tasks 4,5,8,10,11,12 — same fields throughout
rerank_results returns list[tuple[str, str, float]] — consumed correctly in Task 10
get_cache() returns the cache singleton — called with await consistently
deduplicate_results defined in Task 10, imported by Task 12 ✓

76 KiB Raw Blame History Unescape Escape

Axon Search API Implementation Plan

Global Constraints

File Map

Task 1: Project Scaffold

Task 2: Pydantic Models

Task 3: Cache Service

Task 4: Brave Search Source

Task 5: DDG and Serper Sources

Task 6: Scraper Service

Task 7: Reranker Service

Task 8: Synthesizer Service

Task 9: Auth Middleware

Task 10: Search Router

Task 11: Extract Router

Task 12: Research Router

Task 13: Main App + Health Endpoint

Task 14: README

2. Configure environment

3. Start Redis (optional but recommended)

4. Run the server

Endpoints

POST /search

POST /extract

POST /research

GET /health

Configuration

Source Fallback Logic

Performance Notes

Self-Review Against Spec

Spec Coverage

Placeholder Scan

Type Consistency

76 KiB

Raw Blame History