This commit is contained in:
vignesh 2026-03-06 16:39:58 +05:30
commit eeaaa9252e
37 changed files with 3839 additions and 0 deletions

140
.gitignore vendored Normal file
View file

@ -0,0 +1,140 @@
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
pip-wheel-metadata/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
.python-version
# pipenv
Pipfile.lock
# PEP 582
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
# IDE
.vscode/
.idea/
*.swp
*.swo
*~
# OS
.DS_Store
Thumbs.db
# LLM
.claude
.copilot
.gpt
.chatgpt

31
Dockerfile Normal file
View file

@ -0,0 +1,31 @@
FROM python:3.12-slim
WORKDIR /app
# System dependencies for readability-lxml
RUN apt-get update && apt-get install -y --no-install-recommends \
libxml2-dev libxslt-dev gcc \
&& rm -rf /var/lib/apt/lists/*
# Install Python dependencies
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
# Pre-download FlashRank model to avoid first-request delay
RUN python -c "from flashrank import Ranker; Ranker(model_name='ms-marco-MiniLM-L-12-v2')" || true
# Copy application
COPY app/ ./app/
COPY config.yaml .
EXPOSE 8000
# Production: gunicorn with uvicorn workers (uvloop + httptools auto-enabled)
CMD ["gunicorn", "app.main:app", \
"--worker-class", "uvicorn.workers.UvicornWorker", \
"--workers", "4", \
"--bind", "0.0.0.0:8000", \
"--timeout", "60", \
"--graceful-timeout", "30", \
"--keep-alive", "5", \
"--access-logfile", "-"]

422
README.md Normal file
View file

@ -0,0 +1,422 @@
# OrioSearch
Self-hosted, Tavily-compatible web search and content extraction API. Drop-in replacement for Tavily with full control over your search infrastructure.
Built on **SearXNG** (meta-search) + **FastAPI** with async concurrency, Redis caching, circuit breakers, result reranking, SSE streaming, and multi-tier content extraction.
---
## Features
- **Tavily-compatible API** — `/search`, `/extract`, and `/tool-schema` endpoints match Tavily's interface
- **LLM-ready** — `/tool-schema` returns OpenAI function-calling definitions for `web_search` and `web_extract`
- **SearXNG backend** — aggregates 70+ search engines with automatic fallback to DuckDuckGo
- **Content extraction** — multi-tier pipeline: trafilatura (F1: 0.958) with readability-lxml fallback
- **Redis caching** — pipeline-batched lookups, configurable TTLs, stale-cache graceful degradation
- **Result reranking** — FlashRank ONNX model (~4MB, CPU-only, no PyTorch dependency)
- **SSE streaming** — `POST /search/stream` for real-time results via Server-Sent Events
- **Rate limiting** — per-route Redis-backed limits via slowapi, keyed by API key or IP
- **API key auth** — Bearer token authentication with timing-safe comparison
- **Circuit breakers** — automatic failure detection on search backend and extraction
- **Structured logging** — JSON-formatted structlog with request ID correlation
- **Per-domain rate limiting** — LRU-bounded semaphores prevent aggressive crawling
- **Concurrent extraction** — configurable global + per-domain concurrency limits
- **Rotating User-Agents** — 10 browser-like UAs to reduce extraction blocks
- **Proxy support** — optional HTTP/SOCKS proxy for all outbound requests
- **Production deployment** — Gunicorn + UvicornWorker, 4 workers, Docker health checks
---
## Quick Start
### Prerequisites
- Docker and Docker Compose
### Run
```bash
docker compose up --build
```
This starts three services:
| Service | Port | Description |
|---------|------|-------------|
| `orio-search-api` | 8000 | OrioSearch API |
| `orio-search-searxng` | 8080 | SearXNG meta-search engine |
| `orio-search-redis` | 6379 | Redis cache |
### Verify
```bash
curl http://localhost:8000/health
# {"status":"ok","service":"orio-search"}
```
---
## API Reference
### `POST /search`
Search the web and return relevant results.
```json
{
"query": "python async programming",
"search_depth": "basic",
"topic": "general",
"max_results": 5,
"include_images": false,
"include_raw_content": false,
"time_range": "week",
"include_domains": ["docs.python.org"],
"exclude_domains": ["pinterest.com"]
}
```
| Field | Type | Default | Description |
|-------|------|---------|-------------|
| `query` | string | *required* | Search query |
| `search_depth` | `basic` \| `advanced` | `basic` | `advanced` extracts full page content |
| `topic` | `general` \| `news` | `general` | Search category |
| `max_results` | 1-20 | 5 | Number of results |
| `include_images` | boolean | false | Include image results |
| `include_raw_content` | boolean | false | Include extracted page content |
| `time_range` | `day` \| `week` \| `month` \| `year` | null | Time filter |
| `include_domains` | string[] | [] | Whitelist domains |
| `exclude_domains` | string[] | [] | Blacklist domains |
**Response:**
```json
{
"query": "python async programming",
"results": [
{
"title": "Async IO in Python",
"url": "https://docs.python.org/3/library/asyncio.html",
"content": "Snippet of the page content...",
"score": 0.95,
"raw_content": null
}
],
"images": [
{
"url": "https://example.com/image.jpg",
"description": "Python async diagram"
}
],
"response_time": 1.234
}
```
### `POST /search/stream`
Same parameters as `/search`. Returns Server-Sent Events:
```
event: result
data: {"title": "...", "url": "...", "content": "...", "score": 0.95}
event: image
data: {"url": "...", "description": "..."}
event: extraction
data: {"url": "...", "raw_content": "..."}
event: done
data: {"total_results": 5, "response_time": 2.1}
```
### `POST /extract`
Extract clean content from URLs.
```json
{
"urls": ["https://example.com/article"],
"format": "markdown"
}
```
| Field | Type | Default | Description |
|-------|------|---------|-------------|
| `urls` | string[] | *required* | 1-20 URLs to extract |
| `format` | `markdown` \| `text` | `markdown` | Output format |
**Response:**
```json
{
"results": [
{
"url": "https://example.com/article",
"raw_content": "# Article Title\n\nSource: https://example.com/article\n\n---\n\nExtracted content..."
}
],
"failed_results": [
{
"url": "https://example.com/404",
"error": "HTTP 404"
}
],
"response_time": 2.5
}
```
### `GET /tool-schema`
Returns OpenAI-compatible function definitions for `web_search` and `web_extract`. Use this to register OrioSearch as a tool with any LLM.
### `GET /health`
Returns `{"status": "ok", "service": "orio-search"}`.
---
## Configuration
All configuration lives in `config.yaml`. Override the path with the `ORIO_SEARCH_CONFIG` environment variable.
### Search Backend
```yaml
search:
backend: "searxng" # "searxng" | "duckduckgo"
searxng_url: "http://searxng:8080"
```
When `backend_fallback: true` (default), OrioSearch automatically falls back to DuckDuckGo if SearXNG is unavailable.
### Authentication
Disabled by default. Enable to require Bearer tokens:
```yaml
auth:
enabled: true
api_keys:
- "your-secret-key-1"
- "your-secret-key-2"
```
```bash
curl -H "Authorization: Bearer your-secret-key-1" \
-X POST http://localhost:8000/search \
-d '{"query": "test"}'
```
`/health` and `/tool-schema` are always unauthenticated.
### Rate Limiting
```yaml
rate_limit:
enabled: true
search_rate: "30/minute"
extract_rate: "30/minute"
default_rate: "60/minute"
```
When auth is enabled, limits are per API key. Otherwise, per IP.
### Reranking
```yaml
rerank:
enabled: true
model: "ms-marco-MiniLM-L-12-v2"
top_k: 5
```
Uses FlashRank — a ~4MB ONNX model that runs on CPU with no PyTorch dependency. Reranks search results by semantic relevance to the query.
### Caching
```yaml
cache:
enabled: true
redis_url: "redis://redis:6379"
search_ttl: 3600 # 1 hour
extract_ttl: 86400 # 24 hours
```
### Resilience
```yaml
resilience:
circuit_breaker_failure_threshold: 5
circuit_breaker_recovery_timeout: 30
retry_max_attempts: 3
retry_backoff_base: 0.5
request_timeout: 30
backend_fallback: true
```
### Proxy
```yaml
proxy:
enabled: true
url: "socks5://proxy:1080"
```
### Extraction
```yaml
extraction:
max_concurrent: 5 # Global concurrency
domain_concurrency: 2 # Per-domain concurrency
timeout: 10 # Per-URL timeout (seconds)
max_content_length: 50000 # Truncate content after this
domain_semaphore_max_size: 1000 # LRU cache for domain semaphores
```
### CORS
```yaml
cors:
allow_origins: ["*"]
```
### Logging
```yaml
logging:
format: "json" # "json" | "console"
level: "INFO"
```
---
## Using with LLMs
### Get the tool schema
```bash
curl http://localhost:8000/tool-schema
```
This returns OpenAI-compatible function definitions. Pass them to any LLM that supports tool/function calling.
### Example: OpenAI-style tool registration
```python
import requests
# Fetch tool definitions from OrioSearch
schema = requests.get("http://localhost:8000/tool-schema").json()
# Pass to your LLM as tools
response = client.chat.completions.create(
model="your-model",
messages=[{"role": "user", "content": "Find recent news about AI"}],
tools=schema["tools"],
)
# When the LLM calls web_search, forward to OrioSearch
result = requests.post(
"http://localhost:8000/search",
json={"query": "recent AI news", "topic": "news", "max_results": 5},
)
```
---
## Architecture
```
Client
|
v
OrioSearch API (FastAPI + Gunicorn, 4 workers)
|
|-- /search ---------> SearXNG (70+ engines)
| \--> DuckDuckGo (fallback)
|
|-- /extract --------> trafilatura
| \--> readability-lxml (fallback)
|
|-- /search/stream ---> SSE (real-time results)
|
\-- Redis (cache, rate limiting)
```
### Concurrency Model
- **Separate HTTP connection pools**: search (10 connections) and extraction (50 connections)
- **Per-domain semaphores**: LRU-bounded, prevent aggressive crawling of any single domain
- **Global extraction semaphore**: caps total concurrent extractions
- **DuckDuckGo thread pool**: dedicated `ThreadPoolExecutor(4)` for sync DDG calls
- **Redis pipeline batching**: batch extract cache reads/writes into single round-trips
### Resilience
- **Circuit breakers** on search backend and content extraction — auto-open after 5 failures, recover after 30s
- **Exponential backoff retries** for 429/502/503/504 responses
- **Backend fallback**: SearXNG down → automatic DuckDuckGo fallback
- **Graceful degradation**: backend failure → serve stale cached results → HTTP 503
- **Request-level timeouts**: `asyncio.wait_for` wraps every request, returns HTTP 504 on timeout
---
## Development
### Run without Docker
```bash
pip install -r requirements.txt
# Start SearXNG and Redis separately, then:
ORIO_SEARCH_CONFIG=config.yaml uvicorn app.main:app --reload
```
### Run Tests
```bash
pip install pytest pytest-asyncio
pytest tests/ -v
```
---
## Project Structure
```
.
├── app/
│ ├── main.py # FastAPI app, lifespan, middleware
│ ├── config.py # Pydantic config models, YAML loader
│ ├── auth.py # Bearer token authentication
│ ├── rate_limit.py # slowapi limiter instance
│ ├── logging_setup.py # structlog configuration
│ ├── middleware.py # Request ID + timing middleware
│ ├── models/
│ │ └── schemas.py # Pydantic request/response models
│ ├── routers/
│ │ ├── search.py # POST /search
│ │ ├── extract.py # POST /extract
│ │ └── search_stream.py # POST /search/stream (SSE)
│ └── services/
│ ├── search_backend.py # SearXNG + DuckDuckGo + fallback
│ ├── extractor.py # Multi-tier content extraction
│ ├── cache.py # Redis cache with batch ops
│ ├── reranker.py # FlashRank reranking
│ └── resilience.py # Circuit breakers + retry
├── tests/ # pytest test suite
├── searxng/
│ └── settings.yml # SearXNG configuration
├── config.yaml # App configuration
├── docker-compose.yml # 3-service stack
├── Dockerfile # Production image
├── requirements.txt # Python dependencies
└── test.http # VS Code REST Client tests
```
---
## License
MIT

0
app/__init__.py Normal file
View file

25
app/auth.py Normal file
View file

@ -0,0 +1,25 @@
from __future__ import annotations
import secrets
from fastapi import Depends, HTTPException, Security
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from app.config import settings
security = HTTPBearer(auto_error=False)
async def verify_api_key(
credentials: HTTPAuthorizationCredentials | None = Security(security),
) -> str | None:
if not settings.auth.enabled:
return None
if credentials is None:
raise HTTPException(401, "Missing API key. Provide Authorization: Bearer <key>")
if not any(
secrets.compare_digest(credentials.credentials, key)
for key in settings.auth.api_keys
):
raise HTTPException(403, "Invalid API key")
return credentials.credentials

115
app/config.py Normal file
View file

@ -0,0 +1,115 @@
from __future__ import annotations
import os
from pathlib import Path
from typing import Any
import yaml
from pydantic import BaseModel
class ServerConfig(BaseModel):
host: str = "0.0.0.0"
port: int = 8000
class SearchConfig(BaseModel):
backend: str = "searxng"
searxng_url: str = "http://searxng:8080"
google_api_key: str = ""
google_cx: str = ""
class ExtractionConfig(BaseModel):
max_concurrent: int = 5
timeout: int = 10
max_content_length: int = 50000
domain_concurrency: int = 2
domain_semaphore_max_size: int = 1000
user_agents: list[str] = [
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
"Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
"Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:133.0) Gecko/20100101 Firefox/133.0",
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/18.2 Safari/605.1.15",
]
class CacheConfig(BaseModel):
enabled: bool = True
redis_url: str = "redis://redis:6379"
search_ttl: int = 3600
extract_ttl: int = 86400
class ProxyConfig(BaseModel):
enabled: bool = False
url: str = ""
class AuthConfig(BaseModel):
enabled: bool = False
api_keys: list[str] = []
class RateLimitConfig(BaseModel):
enabled: bool = False
default_rate: str = "60/minute"
search_rate: str = "30/minute"
extract_rate: str = "30/minute"
class RerankConfig(BaseModel):
enabled: bool = False
model: str = "ms-marco-MiniLM-L-12-v2"
top_k: int = 5
class ResilienceConfig(BaseModel):
circuit_breaker_failure_threshold: int = 5
circuit_breaker_recovery_timeout: int = 30
retry_max_attempts: int = 3
retry_backoff_base: float = 0.5
retry_on_status_codes: list[int] = [429, 503, 502, 504]
request_timeout: int = 30
backend_fallback: bool = True
class LoggingConfig(BaseModel):
format: str = "json"
level: str = "INFO"
class CorsConfig(BaseModel):
allow_origins: list[str] = ["*"]
class AppConfig(BaseModel):
server: ServerConfig = ServerConfig()
search: SearchConfig = SearchConfig()
extraction: ExtractionConfig = ExtractionConfig()
cache: CacheConfig = CacheConfig()
proxy: ProxyConfig = ProxyConfig()
auth: AuthConfig = AuthConfig()
rate_limit: RateLimitConfig = RateLimitConfig()
rerank: RerankConfig = RerankConfig()
resilience: ResilienceConfig = ResilienceConfig()
logging: LoggingConfig = LoggingConfig()
cors: CorsConfig = CorsConfig()
def load_config(config_path: str | None = None) -> AppConfig:
if config_path is None:
config_path = os.environ.get(
"ORIO_SEARCH_CONFIG",
str(Path(__file__).parent.parent / "config.yaml"),
)
path = Path(config_path)
if path.exists():
with open(path) as f:
data: dict[str, Any] = yaml.safe_load(f) or {}
return AppConfig(**data)
return AppConfig()
settings = load_config()

48
app/logging_setup.py Normal file
View file

@ -0,0 +1,48 @@
from __future__ import annotations
import logging
import structlog
from app.config import settings
def setup_logging() -> None:
shared_processors: list[structlog.types.Processor] = [
structlog.contextvars.merge_contextvars,
structlog.stdlib.add_log_level,
structlog.stdlib.add_logger_name,
structlog.processors.TimeStamper(fmt="iso"),
structlog.processors.StackInfoRenderer(),
structlog.processors.format_exc_info,
]
if settings.logging.format == "json":
renderer: structlog.types.Processor = structlog.processors.JSONRenderer()
else:
renderer = structlog.dev.ConsoleRenderer()
structlog.configure(
processors=[
*shared_processors,
structlog.stdlib.ProcessorFormatter.wrap_for_formatter,
],
logger_factory=structlog.stdlib.LoggerFactory(),
wrapper_class=structlog.stdlib.BoundLogger,
cache_logger_on_first_use=True,
)
formatter = structlog.stdlib.ProcessorFormatter(
processors=[
structlog.stdlib.ProcessorFormatter.remove_processors_meta,
renderer,
],
)
handler = logging.StreamHandler()
handler.setFormatter(formatter)
root = logging.getLogger()
root.handlers.clear()
root.addHandler(handler)
root.setLevel(settings.logging.level)

200
app/main.py Normal file
View file

@ -0,0 +1,200 @@
from __future__ import annotations
from contextlib import asynccontextmanager
import httpx
import structlog
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from slowapi import _rate_limit_exceeded_handler
from slowapi.errors import RateLimitExceeded
from app.config import settings
from app.logging_setup import setup_logging
from app.middleware import RequestIDMiddleware, TimingMiddleware
from app.rate_limit import limiter
from app.routers import extract, search, search_stream
from app.services.cache import CacheService
from app.services.extractor import ContentExtractor
from app.services.reranker import RerankerService
from app.services.search_backend import create_search_backend
setup_logging()
logger = structlog.get_logger(__name__)
TOOL_SCHEMA = {
"tools": [
{
"type": "function",
"function": {
"name": "web_search",
"description": "Search the web and return relevant results with optional content extraction.",
"parameters": {
"type": "object",
"required": ["query"],
"properties": {
"query": {"type": "string", "description": "The search query"},
"search_depth": {
"type": "string",
"enum": ["basic", "advanced"],
"default": "basic",
"description": "basic = snippets only, advanced = full content extraction",
},
"topic": {
"type": "string",
"enum": ["general", "news"],
"default": "general",
},
"max_results": {
"type": "integer",
"minimum": 1,
"maximum": 20,
"default": 5,
},
"include_raw_content": {
"type": "boolean",
"default": False,
"description": "Include full extracted page content",
},
"include_images": {
"type": "boolean",
"default": False,
"description": "Include image search results",
},
"include_domains": {
"type": "array",
"items": {"type": "string"},
"description": "Only include results from these domains",
},
"exclude_domains": {
"type": "array",
"items": {"type": "string"},
"description": "Exclude results from these domains",
},
"time_range": {
"type": "string",
"enum": ["day", "week", "month", "year"],
"description": "Filter by time range",
},
},
},
},
},
{
"type": "function",
"function": {
"name": "web_extract",
"description": "Extract clean content from one or more URLs.",
"parameters": {
"type": "object",
"required": ["urls"],
"properties": {
"urls": {
"type": "array",
"items": {"type": "string"},
"minItems": 1,
"maxItems": 20,
"description": "URLs to extract content from",
},
"format": {
"type": "string",
"enum": ["markdown", "text"],
"default": "markdown",
},
},
},
},
},
]
}
@asynccontextmanager
async def lifespan(app: FastAPI):
logger.info("starting", service="OrioSearch")
proxy_url = settings.proxy.url if settings.proxy.enabled else None
# Separate HTTP clients for search vs extraction
app.state.search_http_client = httpx.AsyncClient(
limits=httpx.Limits(max_connections=10, max_keepalive_connections=5),
timeout=httpx.Timeout(15.0),
proxy=proxy_url,
)
app.state.extract_http_client = httpx.AsyncClient(
limits=httpx.Limits(max_connections=50, max_keepalive_connections=20),
timeout=httpx.Timeout(settings.extraction.timeout + 5),
proxy=proxy_url,
)
# Cache
app.state.cache = CacheService(settings)
await app.state.cache.connect()
# Search backend (with fallback if configured)
app.state.search_backend = create_search_backend(settings, app.state.search_http_client)
# Content extractor
app.state.extractor = ContentExtractor(settings, app.state.extract_http_client)
# Reranker
app.state.reranker = RerankerService()
app.state.reranker.initialize()
logger.info(
"ready",
service="OrioSearch",
backend=settings.search.backend,
cache="enabled" if app.state.cache.enabled else "disabled",
auth="enabled" if settings.auth.enabled else "disabled",
rate_limit="enabled" if settings.rate_limit.enabled else "disabled",
rerank="enabled" if settings.rerank.enabled else "disabled",
fallback="enabled" if settings.resilience.backend_fallback else "disabled",
host=settings.server.host,
port=settings.server.port,
)
yield
await app.state.search_http_client.aclose()
await app.state.extract_http_client.aclose()
await app.state.cache.close()
logger.info("shutdown", service="OrioSearch")
app = FastAPI(
title="OrioSearch",
description="OrioSearch — self-hosted Tavily-compatible web search and content extraction API",
version="2.0.0",
lifespan=lifespan,
)
# Middleware (added in reverse order — Starlette processes bottom-up)
app.add_middleware(
CORSMiddleware,
allow_origins=settings.cors.allow_origins,
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
app.add_middleware(TimingMiddleware)
app.add_middleware(RequestIDMiddleware)
# Rate limiting
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
# Routers
app.include_router(search.router)
app.include_router(extract.router)
app.include_router(search_stream.router)
@app.get("/health")
async def health():
return {"status": "ok", "service": "orio-search"}
@app.get("/tool-schema")
async def tool_schema():
return TOOL_SCHEMA

36
app/middleware.py Normal file
View file

@ -0,0 +1,36 @@
from __future__ import annotations
import time
import uuid
import structlog
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import Response
logger = structlog.get_logger(__name__)
class RequestIDMiddleware(BaseHTTPMiddleware):
async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
request_id = request.headers.get("X-Request-ID", str(uuid.uuid4()))
structlog.contextvars.clear_contextvars()
structlog.contextvars.bind_contextvars(request_id=request_id)
response = await call_next(request)
response.headers["X-Request-ID"] = request_id
return response
class TimingMiddleware(BaseHTTPMiddleware):
async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
start = time.perf_counter()
response = await call_next(request)
duration_ms = round((time.perf_counter() - start) * 1000, 2)
logger.info(
"request_completed",
method=request.method,
path=request.url.path,
status_code=response.status_code,
duration_ms=duration_ms,
)
return response

0
app/models/__init__.py Normal file
View file

97
app/models/schemas.py Normal file
View file

@ -0,0 +1,97 @@
from __future__ import annotations
from enum import Enum
from typing import Optional
from pydantic import BaseModel, Field
# --- Enums ---
class SearchDepth(str, Enum):
basic = "basic"
advanced = "advanced"
class Topic(str, Enum):
general = "general"
news = "news"
class TimeRange(str, Enum):
day = "day"
week = "week"
month = "month"
year = "year"
class ExtractDepth(str, Enum):
basic = "basic"
advanced = "advanced"
class ContentFormat(str, Enum):
markdown = "markdown"
text = "text"
# --- Search ---
class SearchRequest(BaseModel):
query: str = Field(..., description="The search query")
search_depth: SearchDepth = Field(
default=SearchDepth.basic,
description="basic = fast snippets only, advanced = fetch and extract full content",
)
topic: Topic = Field(default=Topic.general, description="Search category")
max_results: int = Field(default=5, ge=1, le=20, description="Number of results to return")
include_answer: bool = Field(default=False, description="Not implemented (reserved for LLM answer generation)")
include_raw_content: bool = Field(default=False, description="Include full extracted page content")
include_images: bool = Field(default=False, description="Include image search results")
include_domains: list[str] = Field(default_factory=list, description="Only include results from these domains")
exclude_domains: list[str] = Field(default_factory=list, description="Exclude results from these domains")
time_range: Optional[TimeRange] = Field(default=None, description="Filter results by time range")
class SearchResult(BaseModel):
title: str
url: str
content: str = Field(description="Short snippet / description")
score: float = Field(default=0.0, description="Relevance score (0-1)")
raw_content: Optional[str] = Field(default=None, description="Full extracted page content")
class ImageResult(BaseModel):
url: str = Field(description="Direct image URL")
description: str = Field(default="", description="Image description or alt text")
class SearchResponse(BaseModel):
query: str
results: list[SearchResult]
images: list[ImageResult] = Field(default_factory=list)
response_time: float = Field(description="Total time in seconds")
# --- Extract ---
class ExtractRequest(BaseModel):
urls: list[str] = Field(..., min_length=1, max_length=20, description="URLs to extract content from")
extract_depth: ExtractDepth = Field(default=ExtractDepth.basic, description="Extraction depth")
format: ContentFormat = Field(default=ContentFormat.markdown, description="Output format")
class ExtractResult(BaseModel):
url: str
raw_content: str = Field(description="Extracted page content")
class FailedResult(BaseModel):
url: str
error: str
class ExtractResponse(BaseModel):
results: list[ExtractResult]
failed_results: list[FailedResult] = Field(default_factory=list)
response_time: float

21
app/rate_limit.py Normal file
View file

@ -0,0 +1,21 @@
from __future__ import annotations
from slowapi import Limiter
from slowapi.util import get_remote_address
from starlette.requests import Request
from app.config import settings
def _get_key(request: Request) -> str:
auth = request.headers.get("Authorization", "")
if auth.startswith("Bearer "):
return auth[7:]
return get_remote_address(request)
limiter = Limiter(
key_func=_get_key,
storage_uri=settings.cache.redis_url if settings.rate_limit.enabled else "memory://",
enabled=settings.rate_limit.enabled,
)

0
app/routers/__init__.py Normal file
View file

75
app/routers/extract.py Normal file
View file

@ -0,0 +1,75 @@
from __future__ import annotations
import asyncio
import time
import structlog
from fastapi import APIRouter, Depends, HTTPException, Request
from app.auth import verify_api_key
from app.config import settings
from app.models.schemas import (
ExtractRequest,
ExtractResponse,
ExtractResult,
FailedResult,
)
from app.rate_limit import limiter
logger = structlog.get_logger(__name__)
router = APIRouter()
@router.post("/extract", response_model=ExtractResponse)
@limiter.limit(settings.rate_limit.extract_rate)
async def extract(
request: Request,
body: ExtractRequest,
api_key: str | None = Depends(verify_api_key),
) -> ExtractResponse:
try:
return await asyncio.wait_for(
_do_extract(request, body),
timeout=settings.resilience.request_timeout,
)
except asyncio.TimeoutError:
raise HTTPException(504, "Request timed out")
async def _do_extract(request: Request, body: ExtractRequest) -> ExtractResponse:
start = time.perf_counter()
cache = request.app.state.cache
extractor = request.app.state.extractor
results: list[ExtractResult] = []
failed: list[FailedResult] = []
# Batch cache lookup (single Redis pipeline round-trip)
cached_map = await cache.get_extract_batch(body.urls)
urls_to_fetch: list[str] = []
for url, content in cached_map.items():
if content is not None:
results.append(ExtractResult(url=url, raw_content=content))
else:
urls_to_fetch.append(url)
# Fetch uncached URLs
if urls_to_fetch:
extractions = await extractor.extract_urls(
urls_to_fetch, output_format=body.format.value,
)
to_cache: list[tuple[str, str]] = []
for extraction in extractions:
if extraction.success:
results.append(ExtractResult(url=extraction.url, raw_content=extraction.content))
to_cache.append((extraction.url, extraction.content))
else:
failed.append(FailedResult(url=extraction.url, error=extraction.error or "Unknown error"))
await cache.set_extract_batch(to_cache)
elapsed = time.perf_counter() - start
return ExtractResponse(
results=results, failed_results=failed,
response_time=round(elapsed, 3),
)

122
app/routers/search.py Normal file
View file

@ -0,0 +1,122 @@
from __future__ import annotations
import asyncio
import hashlib
import json
import time
import structlog
from fastapi import APIRouter, Depends, HTTPException, Request
from app.auth import verify_api_key
from app.config import settings
from app.models.schemas import (
ImageResult,
SearchRequest,
SearchResponse,
SearchResult,
)
from app.rate_limit import limiter
logger = structlog.get_logger(__name__)
router = APIRouter()
def _params_hash(req: SearchRequest) -> str:
data = req.model_dump(exclude={"query"})
return hashlib.sha256(json.dumps(data, sort_keys=True).encode()).hexdigest()[:12]
@router.post("/search", response_model=SearchResponse)
@limiter.limit(settings.rate_limit.search_rate)
async def search(
request: Request,
body: SearchRequest,
api_key: str | None = Depends(verify_api_key),
) -> SearchResponse:
try:
return await asyncio.wait_for(
_do_search(request, body),
timeout=settings.resilience.request_timeout,
)
except asyncio.TimeoutError:
raise HTTPException(504, "Request timed out")
async def _do_search(request: Request, body: SearchRequest) -> SearchResponse:
start = time.perf_counter()
cache = request.app.state.cache
backend = request.app.state.search_backend
extractor = request.app.state.extractor
reranker = request.app.state.reranker
# Check cache
ph = _params_hash(body)
cached = await cache.get_search(body.query, ph)
if cached:
elapsed = time.perf_counter() - start
return SearchResponse(**{**cached, "response_time": round(elapsed, 3)})
# Query search backend with graceful degradation
try:
backend_resp = await backend.search(
body.query,
max_results=body.max_results,
topic=body.topic.value,
time_range=body.time_range.value if body.time_range else None,
include_domains=body.include_domains or None,
exclude_domains=body.exclude_domains or None,
include_images=body.include_images,
)
except Exception as e:
logger.error("search_backend_failed", error=str(e))
stale = await cache.get_search(body.query, ph)
if stale:
elapsed = time.perf_counter() - start
logger.info("serving_stale_cache", query=body.query)
return SearchResponse(**{**stale, "response_time": round(elapsed, 3)})
raise HTTPException(503, "Search service unavailable")
# Build results
results: list[SearchResult] = []
for raw in backend_resp.results:
results.append(
SearchResult(
title=raw.title, url=raw.url, content=raw.snippet,
score=raw.score, raw_content=None,
)
)
# Rerank results
if reranker and settings.rerank.enabled and results:
result_dicts = [r.model_dump() for r in results]
reranked = reranker.rerank(body.query, result_dicts, top_k=body.max_results)
results = [SearchResult(**r) for r in reranked]
# Advanced depth: fetch and extract content
if body.search_depth.value == "advanced" or body.include_raw_content:
urls = [r.url for r in results]
extractions = await extractor.extract_urls(urls)
to_cache: list[tuple[str, str]] = []
for result, extraction in zip(results, extractions):
if extraction.success:
result.raw_content = extraction.content
to_cache.append((result.url, extraction.content))
await cache.set_extract_batch(to_cache)
# Images
images = [ImageResult(url=img.url, description=img.description) for img in backend_resp.images]
elapsed = time.perf_counter() - start
response = SearchResponse(
query=body.query, results=results, images=images,
response_time=round(elapsed, 3),
)
# Cache response
cache_data = response.model_dump()
del cache_data["response_time"]
await cache.set_search(body.query, ph, cache_data)
return response

View file

@ -0,0 +1,81 @@
from __future__ import annotations
import asyncio
import json
import time
import structlog
from fastapi import APIRouter, Depends, Request
from sse_starlette.sse import EventSourceResponse
from app.auth import verify_api_key
from app.config import settings
from app.models.schemas import SearchRequest, SearchResult
from app.rate_limit import limiter
logger = structlog.get_logger(__name__)
router = APIRouter()
@router.post("/search/stream")
@limiter.limit(settings.rate_limit.search_rate)
async def search_stream(
request: Request,
body: SearchRequest,
api_key: str | None = Depends(verify_api_key),
):
async def event_generator():
start = time.perf_counter()
backend = request.app.state.search_backend
extractor = request.app.state.extractor
# Search
try:
backend_resp = await backend.search(
body.query,
max_results=body.max_results,
topic=body.topic.value,
time_range=body.time_range.value if body.time_range else None,
include_domains=body.include_domains or None,
exclude_domains=body.exclude_domains or None,
include_images=body.include_images,
)
except Exception as e:
yield {"event": "error", "data": json.dumps({"error": str(e)})}
return
# Stream each result
results = []
for raw in backend_resp.results:
result = SearchResult(
title=raw.title, url=raw.url, content=raw.snippet,
score=raw.score, raw_content=None,
)
results.append(result)
yield {"event": "result", "data": json.dumps(result.model_dump())}
# Stream images
for img in backend_resp.images:
yield {
"event": "image",
"data": json.dumps({"url": img.url, "description": img.description}),
}
# If advanced, stream extractions as they complete
if body.search_depth.value == "advanced" or body.include_raw_content:
tasks = {
url: asyncio.create_task(extractor.extract_url(r.url))
for r, url in ((r, r.url) for r in results)
}
for url, task in tasks.items():
extraction = await task
if extraction.success:
yield {
"event": "extraction",
"data": json.dumps({"url": url, "raw_content": extraction.content}),
}
elapsed = time.perf_counter() - start
yield {"event": "done", "data": json.dumps({"response_time": round(elapsed, 3)})}
return EventSourceResponse(event_generator())

0
app/services/__init__.py Normal file
View file

121
app/services/cache.py Normal file
View file

@ -0,0 +1,121 @@
from __future__ import annotations
import hashlib
import json
from typing import Any, Optional
import redis.asyncio as redis
import structlog
from app.config import AppConfig
logger = structlog.get_logger(__name__)
class CacheService:
def __init__(self, config: AppConfig) -> None:
self.config = config
self.enabled = config.cache.enabled
self._redis: Optional[redis.Redis] = None
self.search_ttl = config.cache.search_ttl
self.extract_ttl = config.cache.extract_ttl
async def connect(self) -> None:
if not self.enabled:
return
try:
self._redis = redis.from_url(
self.config.cache.redis_url,
decode_responses=True,
)
await self._redis.ping()
logger.info("redis_connected", url=self.config.cache.redis_url)
except Exception as e:
logger.warning("redis_connection_failed", error=str(e))
self.enabled = False
self._redis = None
async def close(self) -> None:
if self._redis:
await self._redis.close()
@staticmethod
def _hash_key(prefix: str, value: str) -> str:
h = hashlib.sha256(value.encode()).hexdigest()[:16]
return f"{prefix}:{h}"
# --- Search cache ---
async def get_search(self, query: str, params_hash: str) -> Optional[dict[str, Any]]:
if not self.enabled or not self._redis:
return None
key = self._hash_key("search", f"{query}|{params_hash}")
try:
data = await self._redis.get(key)
if data:
logger.debug("cache_hit", type="search", query=query)
return json.loads(data)
except Exception as e:
logger.warning("cache_read_error", error=str(e))
return None
async def set_search(self, query: str, params_hash: str, data: dict[str, Any]) -> None:
if not self.enabled or not self._redis:
return
key = self._hash_key("search", f"{query}|{params_hash}")
try:
await self._redis.setex(key, self.search_ttl, json.dumps(data))
except Exception as e:
logger.warning("cache_write_error", error=str(e))
# --- Extract cache (single) ---
async def get_extract(self, url: str) -> Optional[str]:
if not self.enabled or not self._redis:
return None
key = self._hash_key("extract", url)
try:
data = await self._redis.get(key)
if data:
logger.debug("cache_hit", type="extract", url=url)
return data
except Exception as e:
logger.warning("cache_read_error", error=str(e))
return None
async def set_extract(self, url: str, content: str) -> None:
if not self.enabled or not self._redis:
return
key = self._hash_key("extract", url)
try:
await self._redis.setex(key, self.extract_ttl, content)
except Exception as e:
logger.warning("cache_write_error", error=str(e))
# --- Extract cache (batch with pipeline) ---
async def get_extract_batch(self, urls: list[str]) -> dict[str, str | None]:
if not self.enabled or not self._redis:
return {url: None for url in urls}
try:
keys = [self._hash_key("extract", url) for url in urls]
async with self._redis.pipeline(transaction=False) as pipe:
for key in keys:
pipe.get(key)
values = await pipe.execute()
return dict(zip(urls, values))
except Exception as e:
logger.warning("cache_batch_read_error", error=str(e))
return {url: None for url in urls}
async def set_extract_batch(self, pairs: list[tuple[str, str]]) -> None:
if not self.enabled or not self._redis or not pairs:
return
try:
async with self._redis.pipeline(transaction=False) as pipe:
for url, content in pairs:
key = self._hash_key("extract", url)
pipe.setex(key, self.extract_ttl, content)
await pipe.execute()
except Exception as e:
logger.warning("cache_batch_write_error", error=str(e))

172
app/services/extractor.py Normal file
View file

@ -0,0 +1,172 @@
from __future__ import annotations
import asyncio
import random
import re
from collections import OrderedDict
from dataclasses import dataclass
from typing import Optional
from urllib.parse import urlparse
import httpx
import structlog
import trafilatura
from circuitbreaker import CircuitBreakerError
from app.config import AppConfig
from app.services.resilience import extract_circuit, retry_async
logger = structlog.get_logger(__name__)
@dataclass
class ExtractionResult:
url: str
content: Optional[str] = None
error: Optional[str] = None
@property
def success(self) -> bool:
return self.content is not None
class ContentExtractor:
def __init__(self, config: AppConfig, http_client: httpx.AsyncClient) -> None:
self.config = config
self.client = http_client
self._domain_semaphores: OrderedDict[str, asyncio.Semaphore] = OrderedDict()
self._domain_concurrency = config.extraction.domain_concurrency
self._max_domains = config.extraction.domain_semaphore_max_size
self._global_semaphore = asyncio.Semaphore(config.extraction.max_concurrent)
def _get_domain_semaphore(self, url: str) -> asyncio.Semaphore:
domain = urlparse(url).netloc
if domain in self._domain_semaphores:
self._domain_semaphores.move_to_end(domain)
return self._domain_semaphores[domain]
while len(self._domain_semaphores) >= self._max_domains:
self._domain_semaphores.popitem(last=False)
sem = asyncio.Semaphore(self._domain_concurrency)
self._domain_semaphores[domain] = sem
return sem
def _get_headers(self) -> dict[str, str]:
ua = random.choice(self.config.extraction.user_agents)
return {
"User-Agent": ua,
"Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
"Accept-Language": "en-US,en;q=0.9",
"Accept-Encoding": "gzip, deflate, br",
"DNT": "1",
"Connection": "keep-alive",
"Upgrade-Insecure-Requests": "1",
}
async def extract_url(self, url: str, *, output_format: str = "markdown") -> ExtractionResult:
domain_sem = self._get_domain_semaphore(url)
async with self._global_semaphore, domain_sem:
return await self._fetch_and_extract(url, output_format)
async def _fetch_and_extract(self, url: str, output_format: str) -> ExtractionResult:
try:
resp = await self._fetch_url(url)
html = resp.text
content_type = resp.headers.get("content-type", "")
if "text/html" not in content_type and "application/xhtml" not in content_type:
return ExtractionResult(url=url, error=f"Non-HTML content type: {content_type}")
# Tier 1: trafilatura
extracted = self._extract_trafilatura(html, url, output_format)
# Tier 2: readability-lxml fallback
if not extracted:
logger.info("trafilatura_fallback", url=url)
extracted = self._extract_readability(html, output_format)
if not extracted:
return ExtractionResult(url=url, error="No content could be extracted")
max_len = self.config.extraction.max_content_length
if len(extracted) > max_len:
extracted = extracted[:max_len] + "\n\n[Content truncated]"
return ExtractionResult(url=url, content=extracted)
except CircuitBreakerError:
return ExtractionResult(url=url, error="Service temporarily unavailable (circuit open)")
except httpx.TimeoutException:
return ExtractionResult(url=url, error=f"Timeout after {self.config.extraction.timeout}s")
except httpx.HTTPStatusError as e:
return ExtractionResult(url=url, error=f"HTTP {e.response.status_code}")
except Exception as e:
logger.exception("extraction_failed", url=url)
return ExtractionResult(url=url, error=str(e))
async def _fetch_url(self, url: str) -> httpx.Response:
@extract_circuit
async def _do():
return await retry_async(lambda: self._raw_fetch(url))
return await _do()
async def _raw_fetch(self, url: str) -> httpx.Response:
resp = await self.client.get(
url,
headers=self._get_headers(),
timeout=self.config.extraction.timeout,
follow_redirects=True,
)
resp.raise_for_status()
return resp
def _extract_trafilatura(self, html: str, url: str, output_format: str) -> str | None:
extracted = trafilatura.extract(
html,
include_comments=False,
include_tables=True,
favor_recall=True,
url=url,
)
if not extracted:
return None
if output_format == "markdown":
return self._to_markdown(extracted, html, url)
return extracted
def _extract_readability(self, html: str, output_format: str) -> str | None:
try:
from readability import Document as ReadabilityDocument
doc = ReadabilityDocument(html)
content_html = doc.summary()
text = re.sub(r"<[^>]+>", " ", content_html)
text = re.sub(r"\s+", " ", text).strip()
if not text or len(text) < 50:
return None
if output_format == "markdown":
title = doc.short_title() or ""
lines = []
if title:
lines.append(f"# {title}\n")
lines.append("---\n")
lines.append(text)
return "\n".join(lines)
return text
except Exception:
return None
@staticmethod
def _to_markdown(text: str, html: str, url: str) -> str:
title_match = re.search(r"<title[^>]*>(.*?)</title>", html, re.IGNORECASE | re.DOTALL)
title = title_match.group(1).strip() if title_match else ""
lines = []
if title:
lines.append(f"# {title}\n")
lines.append(f"Source: {url}\n")
lines.append("---\n")
lines.append(text)
return "\n".join(lines)
async def extract_urls(self, urls: list[str], *, output_format: str = "markdown") -> list[ExtractionResult]:
tasks = [self.extract_url(url, output_format=output_format) for url in urls]
return await asyncio.gather(*tasks)

52
app/services/reranker.py Normal file
View file

@ -0,0 +1,52 @@
from __future__ import annotations
from typing import Any
import structlog
from app.config import settings
logger = structlog.get_logger(__name__)
class RerankerService:
def __init__(self) -> None:
self._ranker = None
def initialize(self) -> None:
if not settings.rerank.enabled:
return
from flashrank import Ranker
logger.info("loading_rerank_model", model=settings.rerank.model)
self._ranker = Ranker(model_name=settings.rerank.model)
logger.info("rerank_model_loaded")
def rerank(
self,
query: str,
results: list[dict[str, Any]],
top_k: int | None = None,
) -> list[dict[str, Any]]:
if not self._ranker or not results:
return results
from flashrank import RerankRequest
top_k = top_k or settings.rerank.top_k
passages = [
{"id": i, "text": f"{r.get('title', '')} {r.get('content', '')}", "meta": r}
for i, r in enumerate(results)
]
rerank_req = RerankRequest(query=query, passages=passages)
reranked = self._ranker.rerank(rerank_req)
output: list[dict[str, Any]] = []
for item in reranked[:top_k]:
original = item["meta"]
original["score"] = round(item["score"], 4)
output.append(original)
return output

View file

@ -0,0 +1,53 @@
from __future__ import annotations
import asyncio
import httpx
import structlog
from circuitbreaker import circuit
from app.config import settings
logger = structlog.get_logger(__name__)
search_circuit = circuit(
failure_threshold=settings.resilience.circuit_breaker_failure_threshold,
recovery_timeout=settings.resilience.circuit_breaker_recovery_timeout,
expected_exception=Exception,
)
extract_circuit = circuit(
failure_threshold=settings.resilience.circuit_breaker_failure_threshold,
recovery_timeout=settings.resilience.circuit_breaker_recovery_timeout,
expected_exception=Exception,
)
async def retry_async(
coro_factory,
*,
max_attempts: int | None = None,
backoff_base: float | None = None,
retryable_status_codes: list[int] | None = None,
):
_max = max_attempts or settings.resilience.retry_max_attempts
_backoff = backoff_base or settings.resilience.retry_backoff_base
_codes = retryable_status_codes or settings.resilience.retry_on_status_codes
last_exception: Exception | None = None
for attempt in range(1, _max + 1):
try:
return await coro_factory()
except httpx.HTTPStatusError as e:
if e.response.status_code not in _codes:
raise
last_exception = e
except (httpx.ConnectError, httpx.ReadTimeout, httpx.ConnectTimeout) as e:
last_exception = e
if attempt < _max:
delay = _backoff * (2 ** (attempt - 1))
logger.warning("retrying", attempt=attempt, delay=delay, error=str(last_exception))
await asyncio.sleep(delay)
raise last_exception # type: ignore[misc]

View file

@ -0,0 +1,257 @@
from __future__ import annotations
import asyncio
from abc import ABC, abstractmethod
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from typing import Optional
import httpx
import structlog
from circuitbreaker import CircuitBreakerError
from app.config import AppConfig
from app.services.resilience import retry_async, search_circuit
logger = structlog.get_logger(__name__)
@dataclass
class RawSearchResult:
title: str
url: str
snippet: str
score: float = 0.0
@dataclass
class RawImageResult:
url: str
description: str = ""
@dataclass
class BackendSearchResponse:
results: list[RawSearchResult] = field(default_factory=list)
images: list[RawImageResult] = field(default_factory=list)
class SearchBackend(ABC):
@abstractmethod
async def search(
self,
query: str,
*,
max_results: int = 5,
topic: str = "general",
time_range: Optional[str] = None,
include_domains: Optional[list[str]] = None,
exclude_domains: Optional[list[str]] = None,
include_images: bool = False,
) -> BackendSearchResponse:
...
class SearXNGBackend(SearchBackend):
def __init__(self, base_url: str, http_client: httpx.AsyncClient) -> None:
self.base_url = base_url.rstrip("/")
self.client = http_client
async def search(
self,
query: str,
*,
max_results: int = 5,
topic: str = "general",
time_range: Optional[str] = None,
include_domains: Optional[list[str]] = None,
exclude_domains: Optional[list[str]] = None,
include_images: bool = False,
) -> BackendSearchResponse:
effective_query = query
if include_domains:
effective_query += " " + " ".join(f"site:{d}" for d in include_domains)
if exclude_domains:
effective_query += " " + " ".join(f"-site:{d}" for d in exclude_domains)
categories = "news" if topic == "news" else "general"
searxng_time_range = time_range if time_range in ("day", "week", "month", "year") else None
tasks = [self._web_search_with_retry(effective_query, max_results, categories, searxng_time_range)]
if include_images:
tasks.append(self._image_search_with_retry(query, max_results=5))
task_results = await asyncio.gather(*tasks, return_exceptions=True)
web_results: list[RawSearchResult] = []
images: list[RawImageResult] = []
if not isinstance(task_results[0], BaseException):
web_results = task_results[0]
else:
logger.error("web_search_failed", error=str(task_results[0]))
raise task_results[0]
if include_images and len(task_results) > 1:
if not isinstance(task_results[1], BaseException):
images = task_results[1]
else:
logger.warning("image_search_failed", error=str(task_results[1]))
return BackendSearchResponse(results=web_results, images=images)
async def _web_search_with_retry(
self, query: str, max_results: int, categories: str, time_range: Optional[str]
) -> list[RawSearchResult]:
@search_circuit
async def _do():
return await retry_async(
lambda: self._web_search(query, max_results, categories, time_range)
)
return await _do()
async def _image_search_with_retry(self, query: str, max_results: int = 5) -> list[RawImageResult]:
@search_circuit
async def _do():
return await retry_async(lambda: self._image_search(query, max_results))
return await _do()
async def _web_search(
self, query: str, max_results: int, categories: str, time_range: Optional[str]
) -> list[RawSearchResult]:
params: dict = {"q": query, "format": "json", "categories": categories}
if time_range:
params["time_range"] = time_range
resp = await self.client.get(f"{self.base_url}/search", params=params)
resp.raise_for_status()
data = resp.json()
results: list[RawSearchResult] = []
for item in data.get("results", [])[:max_results]:
raw_score = item.get("score", 0)
score = min(1.0, raw_score) if raw_score <= 1.0 else min(1.0, raw_score / 10.0)
results.append(
RawSearchResult(
title=item.get("title", ""),
url=item.get("url", ""),
snippet=item.get("content", ""),
score=round(score, 4),
)
)
return results
async def _image_search(self, query: str, max_results: int = 5) -> list[RawImageResult]:
params = {"q": query, "format": "json", "categories": "images"}
resp = await self.client.get(f"{self.base_url}/search", params=params)
resp.raise_for_status()
data = resp.json()
images: list[RawImageResult] = []
for item in data.get("results", [])[:max_results]:
img_url = item.get("img_src") or item.get("url", "")
if img_url:
images.append(RawImageResult(url=img_url, description=item.get("title", "")))
return images
class DuckDuckGoBackend(SearchBackend):
def __init__(self) -> None:
self._executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="ddg")
async def search(
self,
query: str,
*,
max_results: int = 5,
topic: str = "general",
time_range: Optional[str] = None,
include_domains: Optional[list[str]] = None,
exclude_domains: Optional[list[str]] = None,
include_images: bool = False,
) -> BackendSearchResponse:
effective_query = query
if include_domains:
effective_query += " " + " ".join(f"site:{d}" for d in include_domains)
if exclude_domains:
effective_query += " " + " ".join(f"-site:{d}" for d in exclude_domains)
timelimit = None
if time_range:
timelimit = {"day": "d", "week": "w", "month": "m", "year": "y"}.get(time_range)
loop = asyncio.get_running_loop()
web_results = await loop.run_in_executor(
self._executor,
lambda: self._sync_web_search(effective_query, max_results, timelimit),
)
images: list[RawImageResult] = []
if include_images:
images = await loop.run_in_executor(
self._executor,
lambda: self._sync_image_search(query, max_results=5),
)
return BackendSearchResponse(results=web_results, images=images)
def _sync_web_search(
self, query: str, max_results: int, timelimit: Optional[str]
) -> list[RawSearchResult]:
from duckduckgo_search import DDGS
results: list[RawSearchResult] = []
with DDGS() as ddgs:
raw = list(ddgs.text(query, max_results=max_results, timelimit=timelimit))
for i, item in enumerate(raw):
score = round(1.0 - (i * 0.05), 4)
results.append(
RawSearchResult(
title=item.get("title", ""),
url=item.get("href", ""),
snippet=item.get("body", ""),
score=max(0.1, score),
)
)
return results
def _sync_image_search(self, query: str, max_results: int = 5) -> list[RawImageResult]:
from duckduckgo_search import DDGS
images: list[RawImageResult] = []
with DDGS() as ddgs:
raw = list(ddgs.images(query, max_results=max_results))
for item in raw:
images.append(RawImageResult(url=item.get("image", ""), description=item.get("title", "")))
return images
class FallbackSearchBackend(SearchBackend):
def __init__(self, primary: SearchBackend, fallback: SearchBackend) -> None:
self.primary = primary
self.fallback = fallback
async def search(self, query: str, **kwargs) -> BackendSearchResponse:
try:
return await self.primary.search(query, **kwargs)
except (Exception, CircuitBreakerError) as e:
logger.warning("primary_backend_failed", error=str(e), falling_back=True)
return await self.fallback.search(query, **kwargs)
def create_search_backend(config: AppConfig, http_client: httpx.AsyncClient) -> SearchBackend:
backend_name = config.search.backend.lower()
if backend_name == "searxng":
primary: SearchBackend = SearXNGBackend(config.search.searxng_url, http_client)
elif backend_name == "duckduckgo":
primary = DuckDuckGoBackend()
else:
raise ValueError(f"Unknown search backend: {backend_name}. Supported: searxng, duckduckgo")
if config.resilience.backend_fallback and backend_name == "searxng":
fallback = DuckDuckGoBackend()
return FallbackSearchBackend(primary, fallback)
return primary

70
config.yaml Normal file
View file

@ -0,0 +1,70 @@
server:
host: "0.0.0.0"
port: 8000
search:
backend: "searxng" # "searxng" | "duckduckgo"
searxng_url: "http://searxng:8080"
google_api_key: ""
google_cx: ""
extraction:
max_concurrent: 5
timeout: 10
max_content_length: 50000
domain_concurrency: 2
domain_semaphore_max_size: 1000
user_agents:
- "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36"
- "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36"
- "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:133.0) Gecko/20100101 Firefox/133.0"
- "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/18.2 Safari/605.1.15"
- "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36"
- "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36 Edg/131.0.0.0"
- "Mozilla/5.0 (X11; Linux x86_64; rv:133.0) Gecko/20100101 Firefox/133.0"
- "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:133.0) Gecko/20100101 Firefox/133.0"
- "Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:133.0) Gecko/20100101 Firefox/133.0"
- "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/130.0.0.0 Safari/537.36"
cache:
enabled: true
redis_url: "redis://redis:6379"
search_ttl: 3600
extract_ttl: 86400
proxy:
enabled: false
url: ""
auth:
enabled: false
api_keys: []
# api_keys:
# - "your-secret-key-here"
rate_limit:
enabled: false
default_rate: "60/minute"
search_rate: "30/minute"
extract_rate: "30/minute"
rerank:
enabled: false
model: "ms-marco-MiniLM-L-12-v2"
top_k: 5
resilience:
circuit_breaker_failure_threshold: 5
circuit_breaker_recovery_timeout: 30
retry_max_attempts: 3
retry_backoff_base: 0.5
retry_on_status_codes: [429, 503, 502, 504]
request_timeout: 30
backend_fallback: true
logging:
format: "json" # "json" | "console"
level: "INFO"
cors:
allow_origins: ["*"]

50
docker-compose.yml Normal file
View file

@ -0,0 +1,50 @@
services:
searxng:
image: searxng/searxng:latest
container_name: orio-search-searxng
ports:
- "8080:8080"
volumes:
- ./searxng/settings.yml:/etc/searxng/settings.yml:ro
restart: unless-stopped
healthcheck:
test: ["CMD", "wget", "--no-verbose", "--tries=1", "--spider", "http://localhost:8080/healthz"]
interval: 30s
timeout: 5s
retries: 3
redis:
image: redis:7-alpine
container_name: orio-search-redis
ports:
- "6379:6379"
restart: unless-stopped
healthcheck:
test: ["CMD", "redis-cli", "ping"]
interval: 10s
timeout: 3s
retries: 5
deploy:
resources:
limits:
memory: 256M
api:
build: .
container_name: orio-search-api
ports:
- "8000:8000"
volumes:
- ./config.yaml:/app/config.yaml:ro
depends_on:
searxng:
condition: service_healthy
redis:
condition: service_healthy
environment:
- ORIO_SEARCH_CONFIG=/app/config.yaml
restart: unless-stopped
deploy:
resources:
limits:
memory: 1G

387
quick_test.http Normal file
View file

@ -0,0 +1,387 @@
### ============================================
### OrioSearch API - REST Client Tests
### ============================================
### Requires VS Code extension: REST Client (humao.rest-client)
### Start the stack first: docker compose up --build
@baseUrl = http://localhost:8000
@apiKey = your-secret-key-here
### ============================================
### Health & Info
### ============================================
### Health Check
GET {{baseUrl}}/health
###
### OpenAPI Docs (Swagger)
GET {{baseUrl}}/docs
###
### LLM Tool Schema (OpenAI function-calling format)
GET {{baseUrl}}/tool-schema
###
### ============================================
### SEARCH - Basic
### ============================================
### Basic search (fast, snippets only)
POST {{baseUrl}}/search
Content-Type: application/json
{
"query": "python fastapi tutorial",
"max_results": 5
}
###
### Search with topic = news
POST {{baseUrl}}/search
Content-Type: application/json
{
"query": "artificial intelligence",
"topic": "news",
"max_results": 5
}
###
### Search with time range filter
POST {{baseUrl}}/search
Content-Type: application/json
{
"query": "latest python release",
"time_range": "week",
"max_results": 3
}
###
### ============================================
### SEARCH - Advanced (with content extraction)
### ============================================
### Advanced search - fetches and extracts full page content
POST {{baseUrl}}/search
Content-Type: application/json
{
"query": "how to use docker compose",
"search_depth": "advanced",
"max_results": 3
}
###
### Search with include_raw_content
POST {{baseUrl}}/search
Content-Type: application/json
{
"query": "redis caching best practices",
"include_raw_content": true,
"max_results": 3
}
###
### ============================================
### SEARCH - Images
### ============================================
### Search with images
POST {{baseUrl}}/search
Content-Type: application/json
{
"query": "cute cats",
"include_images": true,
"max_results": 5
}
###
### Advanced search with images
POST {{baseUrl}}/search
Content-Type: application/json
{
"query": "machine learning architecture diagrams",
"search_depth": "advanced",
"include_images": true,
"max_results": 3
}
###
### ============================================
### SEARCH - Domain Filters
### ============================================
### Include only specific domains
POST {{baseUrl}}/search
Content-Type: application/json
{
"query": "python async",
"include_domains": ["docs.python.org", "realpython.com"],
"max_results": 5
}
###
### Exclude specific domains
POST {{baseUrl}}/search
Content-Type: application/json
{
"query": "javascript frameworks 2025",
"exclude_domains": ["medium.com", "w3schools.com"],
"max_results": 5
}
###
### ============================================
### SEARCH - Full featured
### ============================================
### Everything combined
POST {{baseUrl}}/search
Content-Type: application/json
{
"query": "self hosted search engine",
"search_depth": "advanced",
"topic": "general",
"max_results": 5,
"include_raw_content": true,
"include_images": true,
"exclude_domains": ["pinterest.com"],
"time_range": "month"
}
###
### ============================================
### SEARCH - SSE Streaming
### ============================================
### Stream search results (Server-Sent Events)
POST {{baseUrl}}/search/stream
Content-Type: application/json
Accept: text/event-stream
{
"query": "python async programming",
"max_results": 3
}
###
### Stream advanced search with extraction
POST {{baseUrl}}/search/stream
Content-Type: application/json
Accept: text/event-stream
{
"query": "fastapi best practices",
"search_depth": "advanced",
"max_results": 3
}
###
### ============================================
### EXTRACT - Content extraction
### ============================================
### Extract single URL
POST {{baseUrl}}/extract
Content-Type: application/json
{
"urls": ["https://example.com"]
}
###
### Extract multiple URLs (tests Redis pipeline batching)
POST {{baseUrl}}/extract
Content-Type: application/json
{
"urls": [
"https://docs.python.org/3/library/asyncio.html",
"https://fastapi.tiangolo.com/",
"https://httpd.apache.org/"
],
"format": "markdown"
}
###
### Extract with text format
POST {{baseUrl}}/extract
Content-Type: application/json
{
"urls": [
"https://en.wikipedia.org/wiki/Web_scraping"
],
"format": "text"
}
###
### Extract - test error handling (bad URL)
POST {{baseUrl}}/extract
Content-Type: application/json
{
"urls": [
"https://example.com",
"https://thisdomaindoesnotexist12345.com/page"
]
}
###
### ============================================
### AUTH - API Key Authentication
### ============================================
### (Enable auth in config.yaml first:
### auth:
### enabled: true
### api_keys: ["your-secret-key-here"]
### )
### Search with API key (should succeed when auth enabled)
POST {{baseUrl}}/search
Content-Type: application/json
Authorization: Bearer {{apiKey}}
{
"query": "test auth",
"max_results": 3
}
###
### Search without API key (should return 401 when auth enabled)
POST {{baseUrl}}/search
Content-Type: application/json
{
"query": "test no auth",
"max_results": 3
}
###
### Search with invalid API key (should return 403 when auth enabled)
POST {{baseUrl}}/search
Content-Type: application/json
Authorization: Bearer invalid-key-here
{
"query": "test bad auth",
"max_results": 3
}
###
### ============================================
### Edge cases & error handling
### ============================================
### Missing query field (should return 422)
POST {{baseUrl}}/search
Content-Type: application/json
{
"max_results": 5
}
###
### max_results out of range (should return 422)
POST {{baseUrl}}/search
Content-Type: application/json
{
"query": "test",
"max_results": 50
}
###
### Extract with no URLs (should return 422)
POST {{baseUrl}}/extract
Content-Type: application/json
{
"urls": []
}
###
### Extract with too many URLs (should return 422, max 20)
POST {{baseUrl}}/extract
Content-Type: application/json
{
"urls": [
"https://example.com/1", "https://example.com/2", "https://example.com/3",
"https://example.com/4", "https://example.com/5", "https://example.com/6",
"https://example.com/7", "https://example.com/8", "https://example.com/9",
"https://example.com/10", "https://example.com/11", "https://example.com/12",
"https://example.com/13", "https://example.com/14", "https://example.com/15",
"https://example.com/16", "https://example.com/17", "https://example.com/18",
"https://example.com/19", "https://example.com/20", "https://example.com/21"
]
}
###
### ============================================
### Verify X-Request-ID header
### ============================================
### Check that response includes X-Request-ID
POST {{baseUrl}}/search
Content-Type: application/json
{
"query": "request id test",
"max_results": 1
}
### Pass custom request ID
POST {{baseUrl}}/search
Content-Type: application/json
X-Request-ID: my-custom-trace-id-123
{
"query": "custom trace id test",
"max_results": 1
}

20
requirements.txt Normal file
View file

@ -0,0 +1,20 @@
fastapi==0.115.6
uvicorn[standard]==0.34.0
httpx==0.28.1
trafilatura==2.0.0
redis[hiredis]==5.2.1
pyyaml==6.0.2
pydantic==2.10.4
pydantic-settings==2.7.1
duckduckgo-search==7.5.3
gunicorn==23.0.0
slowapi==0.1.9
structlog==24.4.0
sse-starlette==2.2.1
flashrank==0.2.9
readability-lxml==0.8.1
circuitbreaker==2.0.0
# Dev / Testing
pytest==8.3.4
pytest-asyncio==0.25.0

13
searxng/settings.yml Normal file
View file

@ -0,0 +1,13 @@
use_default_settings: true
server:
limiter: false
secret_key: "orio-search-secret-change-me"
search:
formats:
- html
- json
general:
enable_metrics: false

0
tests/__init__.py Normal file
View file

127
tests/conftest.py Normal file
View file

@ -0,0 +1,127 @@
from __future__ import annotations
import asyncio
from collections.abc import AsyncGenerator
from typing import Any
from unittest.mock import AsyncMock, MagicMock
import pytest
import pytest_asyncio
from httpx import ASGITransport, AsyncClient
from app.config import AppConfig
from app.services.cache import CacheService
from app.services.extractor import ContentExtractor, ExtractionResult
from app.services.reranker import RerankerService
from app.services.search_backend import (
BackendSearchResponse,
RawImageResult,
RawSearchResult,
SearchBackend,
)
# ---- Fake search backend ----
class FakeSearchBackend(SearchBackend):
def __init__(self, fail: bool = False) -> None:
self.fail = fail
self.last_kwargs: dict = {}
async def search(self, query: str, **kwargs) -> BackendSearchResponse:
self.last_kwargs = {"query": query, **kwargs}
if self.fail:
raise RuntimeError("Backend unavailable")
results = [
RawSearchResult(title=f"Result {i}", url=f"https://example.com/{i}", snippet=f"Snippet {i}", score=round(1.0 - i * 0.1, 2))
for i in range(kwargs.get("max_results", 3))
]
images = []
if kwargs.get("include_images"):
images = [
RawImageResult(url="https://example.com/img1.jpg", description="Test image"),
]
return BackendSearchResponse(results=results, images=images)
# ---- Fake extractor ----
class FakeExtractor:
async def extract_url(self, url: str, *, output_format: str = "markdown") -> ExtractionResult:
if "fail" in url:
return ExtractionResult(url=url, error="Extraction failed")
return ExtractionResult(url=url, content=f"# Extracted content from {url}\n\nFull text here.")
async def extract_urls(self, urls: list[str], *, output_format: str = "markdown") -> list[ExtractionResult]:
return await asyncio.gather(*[self.extract_url(url, output_format=output_format) for url in urls])
# ---- Fake cache ----
class FakeCache:
def __init__(self) -> None:
self._search: dict[str, dict] = {}
self._extract: dict[str, str] = {}
self.enabled = True
async def get_search(self, query: str, params_hash: str) -> dict | None:
return self._search.get(f"{query}|{params_hash}")
async def set_search(self, query: str, params_hash: str, data: dict) -> None:
self._search[f"{query}|{params_hash}"] = data
async def get_extract(self, url: str) -> str | None:
return self._extract.get(url)
async def set_extract(self, url: str, content: str) -> None:
self._extract[url] = content
async def get_extract_batch(self, urls: list[str]) -> dict[str, str | None]:
return {url: self._extract.get(url) for url in urls}
async def set_extract_batch(self, pairs: list[tuple[str, str]]) -> None:
for url, content in pairs:
self._extract[url] = content
async def connect(self) -> None:
pass
async def close(self) -> None:
pass
# ---- App fixture ----
@pytest_asyncio.fixture
async def app():
"""Create a test app with faked services."""
# Import here to avoid module-level config loading issues
from app.main import app as _app
_app.state.search_backend = FakeSearchBackend()
_app.state.extractor = FakeExtractor()
_app.state.cache = FakeCache()
_app.state.reranker = RerankerService() # disabled by default
yield _app
@pytest_asyncio.fixture
async def client(app) -> AsyncGenerator[AsyncClient, None]:
transport = ASGITransport(app=app)
async with AsyncClient(transport=transport, base_url="http://test") as ac:
yield ac
@pytest_asyncio.fixture
async def app_with_failing_backend(app):
"""App with a search backend that always fails."""
app.state.search_backend = FakeSearchBackend(fail=True)
return app
@pytest_asyncio.fixture
async def client_failing_backend(app_with_failing_backend) -> AsyncGenerator[AsyncClient, None]:
transport = ASGITransport(app=app_with_failing_backend)
async with AsyncClient(transport=transport, base_url="http://test") as ac:
yield ac

141
tests/test_auth.py Normal file
View file

@ -0,0 +1,141 @@
from __future__ import annotations
from unittest.mock import patch
import pytest
from httpx import ASGITransport, AsyncClient
@pytest.mark.asyncio
async def test_auth_disabled_no_key_required(client):
"""When auth is disabled (default), requests work without API key."""
resp = await client.post("/search", json={"query": "test", "max_results": 1})
assert resp.status_code == 200
@pytest.mark.asyncio
async def test_auth_disabled_key_ignored(client):
"""When auth is disabled, providing a key still works."""
resp = await client.post(
"/search",
json={"query": "test", "max_results": 1},
headers={"Authorization": "Bearer some-key"},
)
assert resp.status_code == 200
@pytest.mark.asyncio
async def test_health_no_auth_required(client):
"""/health is always unauthenticated."""
resp = await client.get("/health")
assert resp.status_code == 200
@pytest.mark.asyncio
async def test_tool_schema_no_auth_required(client):
"""/tool-schema is always unauthenticated."""
resp = await client.get("/tool-schema")
assert resp.status_code == 200
@pytest.mark.asyncio
async def test_auth_enabled_missing_key(app):
"""When auth is enabled, missing key returns 401."""
with patch("app.auth.settings") as mock_settings:
mock_settings.auth.enabled = True
mock_settings.auth.api_keys = ["valid-key-123"]
transport = ASGITransport(app=app)
async with AsyncClient(transport=transport, base_url="http://test") as ac:
resp = await ac.post("/search", json={"query": "test", "max_results": 1})
assert resp.status_code == 401
@pytest.mark.asyncio
async def test_auth_enabled_invalid_key(app):
"""When auth is enabled, wrong key returns 403."""
with patch("app.auth.settings") as mock_settings:
mock_settings.auth.enabled = True
mock_settings.auth.api_keys = ["valid-key-123"]
transport = ASGITransport(app=app)
async with AsyncClient(transport=transport, base_url="http://test") as ac:
resp = await ac.post(
"/search",
json={"query": "test", "max_results": 1},
headers={"Authorization": "Bearer wrong-key"},
)
assert resp.status_code == 403
@pytest.mark.asyncio
async def test_auth_enabled_valid_key(app):
"""When auth is enabled, correct key works."""
with patch("app.auth.settings") as mock_settings:
mock_settings.auth.enabled = True
mock_settings.auth.api_keys = ["valid-key-123"]
transport = ASGITransport(app=app)
async with AsyncClient(transport=transport, base_url="http://test") as ac:
resp = await ac.post(
"/search",
json={"query": "test", "max_results": 1},
headers={"Authorization": "Bearer valid-key-123"},
)
assert resp.status_code == 200
@pytest.mark.asyncio
async def test_auth_enabled_extract_endpoint(app):
"""Auth also applies to /extract."""
with patch("app.auth.settings") as mock_settings:
mock_settings.auth.enabled = True
mock_settings.auth.api_keys = ["valid-key-123"]
transport = ASGITransport(app=app)
async with AsyncClient(transport=transport, base_url="http://test") as ac:
# Without key
resp = await ac.post("/extract", json={"urls": ["https://example.com"]})
assert resp.status_code == 401
# With valid key
resp = await ac.post(
"/extract",
json={"urls": ["https://example.com"]},
headers={"Authorization": "Bearer valid-key-123"},
)
assert resp.status_code == 200
@pytest.mark.asyncio
async def test_auth_enabled_stream_endpoint(app):
"""Auth also applies to /search/stream."""
with patch("app.auth.settings") as mock_settings:
mock_settings.auth.enabled = True
mock_settings.auth.api_keys = ["valid-key-123"]
transport = ASGITransport(app=app)
async with AsyncClient(transport=transport, base_url="http://test") as ac:
resp = await ac.post(
"/search/stream",
json={"query": "test", "max_results": 1},
)
assert resp.status_code == 401
@pytest.mark.asyncio
async def test_auth_multiple_keys(app):
"""Multiple API keys should all work."""
with patch("app.auth.settings") as mock_settings:
mock_settings.auth.enabled = True
mock_settings.auth.api_keys = ["key-1", "key-2", "key-3"]
transport = ASGITransport(app=app)
async with AsyncClient(transport=transport, base_url="http://test") as ac:
for key in ["key-1", "key-2", "key-3"]:
resp = await ac.post(
"/search",
json={"query": "test", "max_results": 1},
headers={"Authorization": f"Bearer {key}"},
)
assert resp.status_code == 200, f"Key {key} should be valid"

103
tests/test_config.py Normal file
View file

@ -0,0 +1,103 @@
from __future__ import annotations
import os
import tempfile
import pytest
import yaml
from app.config import AppConfig, load_config
class TestConfig:
def test_default_config(self):
config = AppConfig()
assert config.server.host == "0.0.0.0"
assert config.server.port == 8000
assert config.search.backend == "searxng"
assert config.cache.enabled is True
assert config.auth.enabled is False
assert config.rate_limit.enabled is False
assert config.rerank.enabled is False
assert config.resilience.backend_fallback is True
assert config.logging.format == "json"
assert config.cors.allow_origins == ["*"]
def test_extraction_defaults(self):
config = AppConfig()
assert config.extraction.max_concurrent == 5
assert config.extraction.timeout == 10
assert config.extraction.domain_concurrency == 2
assert config.extraction.domain_semaphore_max_size == 1000
assert len(config.extraction.user_agents) > 0
def test_resilience_defaults(self):
config = AppConfig()
assert config.resilience.circuit_breaker_failure_threshold == 5
assert config.resilience.circuit_breaker_recovery_timeout == 30
assert config.resilience.retry_max_attempts == 3
assert config.resilience.retry_backoff_base == 0.5
assert 429 in config.resilience.retry_on_status_codes
assert 503 in config.resilience.retry_on_status_codes
assert config.resilience.request_timeout == 30
def test_load_from_yaml(self):
data = {
"server": {"port": 9999},
"auth": {"enabled": True, "api_keys": ["key1"]},
"rerank": {"enabled": True, "model": "test-model"},
}
with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
yaml.dump(data, f)
f.flush()
config = load_config(f.name)
os.unlink(f.name)
assert config.server.port == 9999
assert config.auth.enabled is True
assert config.auth.api_keys == ["key1"]
assert config.rerank.enabled is True
assert config.rerank.model == "test-model"
# Defaults preserved for unset
assert config.search.backend == "searxng"
assert config.cache.enabled is True
def test_load_missing_file_returns_defaults(self):
config = load_config("/nonexistent/path/config.yaml")
assert config.server.port == 8000
assert config.search.backend == "searxng"
def test_load_empty_yaml(self):
with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
f.write("")
f.flush()
config = load_config(f.name)
os.unlink(f.name)
assert config.server.port == 8000
def test_rate_limit_config(self):
config = AppConfig(rate_limit={"enabled": True, "search_rate": "10/minute"})
assert config.rate_limit.enabled is True
assert config.rate_limit.search_rate == "10/minute"
assert config.rate_limit.extract_rate == "30/minute" # default
def test_proxy_config(self):
config = AppConfig(proxy={"enabled": True, "url": "http://proxy:8080"})
assert config.proxy.enabled is True
assert config.proxy.url == "http://proxy:8080"
def test_cors_config(self):
config = AppConfig(cors={"allow_origins": ["https://myapp.com"]})
assert config.cors.allow_origins == ["https://myapp.com"]
def test_env_var_config_path(self):
"""ORIO_SEARCH_CONFIG env var should be checked."""
with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
yaml.dump({"server": {"port": 7777}}, f)
f.flush()
os.environ["ORIO_SEARCH_CONFIG"] = f.name
config = load_config()
del os.environ["ORIO_SEARCH_CONFIG"]
os.unlink(f.name)
assert config.server.port == 7777

153
tests/test_extract.py Normal file
View file

@ -0,0 +1,153 @@
from __future__ import annotations
import pytest
# ---- Basic extraction ----
@pytest.mark.asyncio
async def test_extract_single_url(client):
resp = await 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]["url"] == "https://example.com"
assert "Extracted content" in data["results"][0]["raw_content"]
assert data["failed_results"] == []
assert data["response_time"] > 0
@pytest.mark.asyncio
async def test_extract_multiple_urls(client):
urls = ["https://example.com/1", "https://example.com/2", "https://example.com/3"]
resp = await client.post("/extract", json={"urls": urls})
assert resp.status_code == 200
data = resp.json()
assert len(data["results"]) == 3
returned_urls = {r["url"] for r in data["results"]}
assert returned_urls == set(urls)
@pytest.mark.asyncio
async def test_extract_with_text_format(client):
resp = await client.post("/extract", json={
"urls": ["https://example.com"],
"format": "text",
})
assert resp.status_code == 200
assert len(resp.json()["results"]) == 1
@pytest.mark.asyncio
async def test_extract_with_markdown_format(client):
resp = await client.post("/extract", json={
"urls": ["https://example.com"],
"format": "markdown",
})
assert resp.status_code == 200
# ---- Failed extractions ----
@pytest.mark.asyncio
async def test_extract_failed_url(client):
resp = await client.post("/extract", json={
"urls": ["https://fail.example.com/page"],
})
assert resp.status_code == 200
data = resp.json()
assert len(data["results"]) == 0
assert len(data["failed_results"]) == 1
assert data["failed_results"][0]["url"] == "https://fail.example.com/page"
assert data["failed_results"][0]["error"] != ""
@pytest.mark.asyncio
async def test_extract_mixed_success_and_failure(client):
resp = await client.post("/extract", json={
"urls": ["https://example.com/ok", "https://fail.example.com/bad"],
})
assert resp.status_code == 200
data = resp.json()
assert len(data["results"]) == 1
assert len(data["failed_results"]) == 1
assert data["results"][0]["url"] == "https://example.com/ok"
assert data["failed_results"][0]["url"] == "https://fail.example.com/bad"
# ---- Cache behavior ----
@pytest.mark.asyncio
async def test_extract_caches_results(client, app):
url = "https://example.com/cached"
resp1 = await client.post("/extract", json={"urls": [url]})
assert resp1.status_code == 200
# Verify it's in cache
cache = app.state.cache
cached = await cache.get_extract(url)
assert cached is not None
assert "Extracted content" in cached
@pytest.mark.asyncio
async def test_extract_serves_from_cache(client, app):
url = "https://example.com/precached"
# Pre-populate cache
await app.state.cache.set_extract(url, "Cached content here")
resp = await client.post("/extract", json={"urls": [url]})
assert resp.status_code == 200
assert resp.json()["results"][0]["raw_content"] == "Cached content here"
@pytest.mark.asyncio
async def test_extract_batch_cache_partial_hit(client, app):
"""Some URLs cached, some not — should fetch only uncached."""
await app.state.cache.set_extract("https://example.com/a", "Cached A")
resp = await client.post("/extract", json={
"urls": ["https://example.com/a", "https://example.com/b"],
})
assert resp.status_code == 200
data = resp.json()
assert len(data["results"]) == 2
urls_to_content = {r["url"]: r["raw_content"] for r in data["results"]}
assert urls_to_content["https://example.com/a"] == "Cached A"
assert "Extracted content" in urls_to_content["https://example.com/b"]
# ---- Validation errors ----
@pytest.mark.asyncio
async def test_extract_empty_urls(client):
resp = await client.post("/extract", json={"urls": []})
assert resp.status_code == 422
@pytest.mark.asyncio
async def test_extract_too_many_urls(client):
urls = [f"https://example.com/{i}" for i in range(21)]
resp = await client.post("/extract", json={"urls": urls})
assert resp.status_code == 422
@pytest.mark.asyncio
async def test_extract_missing_urls(client):
resp = await client.post("/extract", json={})
assert resp.status_code == 422
@pytest.mark.asyncio
async def test_extract_invalid_format(client):
resp = await client.post("/extract", json={
"urls": ["https://example.com"],
"format": "html",
})
assert resp.status_code == 422
@pytest.mark.asyncio
async def test_extract_no_body(client):
resp = await client.post("/extract")
assert resp.status_code == 422

62
tests/test_health.py Normal file
View file

@ -0,0 +1,62 @@
from __future__ import annotations
import pytest
@pytest.mark.asyncio
async def test_health(client):
resp = await client.get("/health")
assert resp.status_code == 200
data = resp.json()
assert data["status"] == "ok"
assert data["service"] == "orio-search"
@pytest.mark.asyncio
async def test_tool_schema(client):
resp = await client.get("/tool-schema")
assert resp.status_code == 200
data = resp.json()
assert "tools" in data
assert len(data["tools"]) == 2
tool_names = {t["function"]["name"] for t in data["tools"]}
assert tool_names == {"web_search", "web_extract"}
@pytest.mark.asyncio
async def test_tool_schema_search_params(client):
resp = await client.get("/tool-schema")
tools = resp.json()["tools"]
search_tool = next(t for t in tools if t["function"]["name"] == "web_search")
props = search_tool["function"]["parameters"]["properties"]
assert "query" in props
assert "search_depth" in props
assert "topic" in props
assert "max_results" in props
assert "include_images" in props
assert "time_range" in props
assert "include_domains" in props
assert "exclude_domains" in props
@pytest.mark.asyncio
async def test_openapi_docs(client):
resp = await client.get("/openapi.json")
assert resp.status_code == 200
data = resp.json()
assert data["info"]["title"] == "OrioSearch"
assert data["info"]["version"] == "2.0.0"
@pytest.mark.asyncio
async def test_request_id_header(client):
resp = await client.get("/health")
assert "x-request-id" in resp.headers
@pytest.mark.asyncio
async def test_custom_request_id(client):
resp = await client.get("/health", headers={"X-Request-ID": "my-trace-123"})
assert resp.headers["x-request-id"] == "my-trace-123"

151
tests/test_schemas.py Normal file
View file

@ -0,0 +1,151 @@
from __future__ import annotations
import pytest
from pydantic import ValidationError
from app.models.schemas import (
ContentFormat,
ExtractRequest,
ExtractResponse,
ExtractResult,
FailedResult,
ImageResult,
SearchDepth,
SearchRequest,
SearchResponse,
SearchResult,
TimeRange,
Topic,
)
class TestSearchRequest:
def test_minimal(self):
req = SearchRequest(query="hello")
assert req.query == "hello"
assert req.search_depth == SearchDepth.basic
assert req.topic == Topic.general
assert req.max_results == 5
assert req.include_images is False
assert req.include_raw_content is False
assert req.time_range is None
assert req.include_domains == []
assert req.exclude_domains == []
def test_full(self):
req = SearchRequest(
query="test",
search_depth="advanced",
topic="news",
max_results=10,
include_images=True,
include_raw_content=True,
time_range="week",
include_domains=["example.com"],
exclude_domains=["spam.com"],
)
assert req.search_depth == SearchDepth.advanced
assert req.topic == Topic.news
assert req.time_range == TimeRange.week
def test_max_results_validation(self):
with pytest.raises(ValidationError):
SearchRequest(query="test", max_results=0)
with pytest.raises(ValidationError):
SearchRequest(query="test", max_results=21)
def test_invalid_depth(self):
with pytest.raises(ValidationError):
SearchRequest(query="test", search_depth="ultra")
def test_invalid_topic(self):
with pytest.raises(ValidationError):
SearchRequest(query="test", topic="sports")
def test_invalid_time_range(self):
with pytest.raises(ValidationError):
SearchRequest(query="test", time_range="decade")
def test_missing_query(self):
with pytest.raises(ValidationError):
SearchRequest()
class TestSearchResponse:
def test_structure(self):
resp = SearchResponse(
query="test",
results=[
SearchResult(title="T", url="https://u.com", content="C", score=0.9),
],
images=[
ImageResult(url="https://img.com/1.jpg", description="desc"),
],
response_time=1.23,
)
assert resp.query == "test"
assert len(resp.results) == 1
assert len(resp.images) == 1
assert resp.response_time == 1.23
def test_empty_results(self):
resp = SearchResponse(query="test", results=[], response_time=0.5)
assert resp.results == []
assert resp.images == []
def test_result_with_raw_content(self):
result = SearchResult(
title="T", url="https://u.com", content="C",
score=0.8, raw_content="Full text here",
)
assert result.raw_content == "Full text here"
def test_result_without_raw_content(self):
result = SearchResult(title="T", url="https://u.com", content="C")
assert result.raw_content is None
assert result.score == 0.0
class TestExtractRequest:
def test_minimal(self):
req = ExtractRequest(urls=["https://example.com"])
assert req.format == ContentFormat.markdown
def test_text_format(self):
req = ExtractRequest(urls=["https://example.com"], format="text")
assert req.format == ContentFormat.text
def test_empty_urls(self):
with pytest.raises(ValidationError):
ExtractRequest(urls=[])
def test_too_many_urls(self):
urls = [f"https://example.com/{i}" for i in range(21)]
with pytest.raises(ValidationError):
ExtractRequest(urls=urls)
def test_max_urls(self):
urls = [f"https://example.com/{i}" for i in range(20)]
req = ExtractRequest(urls=urls)
assert len(req.urls) == 20
def test_invalid_format(self):
with pytest.raises(ValidationError):
ExtractRequest(urls=["https://example.com"], format="html")
class TestExtractResponse:
def test_structure(self):
resp = ExtractResponse(
results=[ExtractResult(url="https://u.com", raw_content="content")],
failed_results=[FailedResult(url="https://bad.com", error="timeout")],
response_time=2.5,
)
assert len(resp.results) == 1
assert len(resp.failed_results) == 1
assert resp.response_time == 2.5
def test_empty_response(self):
resp = ExtractResponse(results=[], response_time=0.1)
assert resp.results == []
assert resp.failed_results == []

214
tests/test_search.py Normal file
View file

@ -0,0 +1,214 @@
from __future__ import annotations
import pytest
# ---- Basic search ----
@pytest.mark.asyncio
async def test_search_basic(client):
resp = await client.post("/search", json={"query": "python fastapi", "max_results": 3})
assert resp.status_code == 200
data = resp.json()
assert data["query"] == "python fastapi"
assert len(data["results"]) == 3
assert "response_time" in data
assert data["response_time"] > 0
@pytest.mark.asyncio
async def test_search_result_fields(client):
resp = await client.post("/search", json={"query": "test", "max_results": 1})
result = resp.json()["results"][0]
assert "title" in result
assert "url" in result
assert "content" in result
assert "score" in result
assert isinstance(result["score"], float)
assert result["raw_content"] is None # basic depth
@pytest.mark.asyncio
async def test_search_default_max_results(client):
resp = await client.post("/search", json={"query": "test"})
assert resp.status_code == 200
assert len(resp.json()["results"]) == 5 # default
@pytest.mark.asyncio
async def test_search_max_results_range(client):
resp = await client.post("/search", json={"query": "test", "max_results": 1})
assert len(resp.json()["results"]) == 1
resp = await client.post("/search", json={"query": "test", "max_results": 20})
assert len(resp.json()["results"]) == 20
# ---- Search with images ----
@pytest.mark.asyncio
async def test_search_with_images(client):
resp = await client.post("/search", json={"query": "cats", "include_images": True})
assert resp.status_code == 200
data = resp.json()
assert len(data["images"]) > 0
assert "url" in data["images"][0]
assert "description" in data["images"][0]
@pytest.mark.asyncio
async def test_search_without_images(client):
resp = await client.post("/search", json={"query": "test", "include_images": False})
assert resp.json()["images"] == []
# ---- Search with advanced depth ----
@pytest.mark.asyncio
async def test_search_advanced_depth(client):
resp = await client.post("/search", json={
"query": "test", "search_depth": "advanced", "max_results": 2,
})
assert resp.status_code == 200
results = resp.json()["results"]
for r in results:
assert r["raw_content"] is not None
assert "Extracted content" in r["raw_content"]
@pytest.mark.asyncio
async def test_search_include_raw_content(client):
resp = await client.post("/search", json={
"query": "test", "include_raw_content": True, "max_results": 1,
})
assert resp.json()["results"][0]["raw_content"] is not None
# ---- Search with filters ----
@pytest.mark.asyncio
async def test_search_topic_news(client, app):
resp = await client.post("/search", json={"query": "ai", "topic": "news", "max_results": 1})
assert resp.status_code == 200
backend = app.state.search_backend
assert backend.last_kwargs["topic"] == "news"
@pytest.mark.asyncio
async def test_search_time_range(client, app):
resp = await client.post("/search", json={"query": "test", "time_range": "week", "max_results": 1})
assert resp.status_code == 200
assert app.state.search_backend.last_kwargs["time_range"] == "week"
@pytest.mark.asyncio
async def test_search_include_domains(client, app):
resp = await client.post("/search", json={
"query": "test", "include_domains": ["example.com"], "max_results": 1,
})
assert resp.status_code == 200
assert app.state.search_backend.last_kwargs["include_domains"] == ["example.com"]
@pytest.mark.asyncio
async def test_search_exclude_domains(client, app):
resp = await client.post("/search", json={
"query": "test", "exclude_domains": ["spam.com"], "max_results": 1,
})
assert resp.status_code == 200
assert app.state.search_backend.last_kwargs["exclude_domains"] == ["spam.com"]
# ---- Search caching ----
@pytest.mark.asyncio
async def test_search_cache_hit(client):
# First request populates cache
resp1 = await client.post("/search", json={"query": "cache test", "max_results": 2})
assert resp1.status_code == 200
# Second identical request should hit cache (faster)
resp2 = await client.post("/search", json={"query": "cache test", "max_results": 2})
assert resp2.status_code == 200
assert resp2.json()["results"] == resp1.json()["results"]
# ---- Search validation errors ----
@pytest.mark.asyncio
async def test_search_missing_query(client):
resp = await client.post("/search", json={"max_results": 5})
assert resp.status_code == 422
@pytest.mark.asyncio
async def test_search_max_results_too_high(client):
resp = await client.post("/search", json={"query": "test", "max_results": 50})
assert resp.status_code == 422
@pytest.mark.asyncio
async def test_search_max_results_too_low(client):
resp = await client.post("/search", json={"query": "test", "max_results": 0})
assert resp.status_code == 422
@pytest.mark.asyncio
async def test_search_invalid_depth(client):
resp = await client.post("/search", json={"query": "test", "search_depth": "ultra"})
assert resp.status_code == 422
@pytest.mark.asyncio
async def test_search_invalid_topic(client):
resp = await client.post("/search", json={"query": "test", "topic": "finance"})
assert resp.status_code == 422
@pytest.mark.asyncio
async def test_search_invalid_time_range(client):
resp = await client.post("/search", json={"query": "test", "time_range": "century"})
assert resp.status_code == 422
@pytest.mark.asyncio
async def test_search_empty_body(client):
resp = await client.post("/search", json={})
assert resp.status_code == 422
@pytest.mark.asyncio
async def test_search_no_body(client):
resp = await client.post("/search")
assert resp.status_code == 422
# ---- Graceful degradation ----
@pytest.mark.asyncio
async def test_search_backend_failure_returns_503(client_failing_backend):
resp = await client_failing_backend.post("/search", json={"query": "test", "max_results": 1})
assert resp.status_code == 503
@pytest.mark.asyncio
async def test_search_backend_failure_serves_stale_cache(app_with_failing_backend):
"""When backend fails but stale cache exists, serve it."""
from httpx import ASGITransport, AsyncClient
from tests.conftest import FakeSearchBackend
# First, populate cache with working backend
app_with_failing_backend.state.search_backend = FakeSearchBackend(fail=False)
transport = ASGITransport(app=app_with_failing_backend)
async with AsyncClient(transport=transport, base_url="http://test") as ac:
resp = await ac.post("/search", json={"query": "stale test", "max_results": 2})
assert resp.status_code == 200
original_results = resp.json()["results"]
# Now break backend and expect stale cache
app_with_failing_backend.state.search_backend = FakeSearchBackend(fail=True)
transport = ASGITransport(app=app_with_failing_backend)
async with AsyncClient(transport=transport, base_url="http://test") as ac:
resp = await ac.post("/search", json={"query": "stale test", "max_results": 2})
assert resp.status_code == 200
assert resp.json()["results"] == original_results

183
tests/test_services.py Normal file
View file

@ -0,0 +1,183 @@
from __future__ import annotations
import asyncio
from collections import OrderedDict
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from app.config import AppConfig
from app.services.cache import CacheService
from app.services.extractor import ContentExtractor, ExtractionResult
from app.services.reranker import RerankerService
from app.services.search_backend import (
BackendSearchResponse,
DuckDuckGoBackend,
FallbackSearchBackend,
RawSearchResult,
SearXNGBackend,
)
# ========== Cache service tests ==========
class TestCacheService:
@pytest.mark.asyncio
async def test_disabled_cache_returns_none(self):
config = AppConfig()
config.cache.enabled = False
cache = CacheService(config)
assert await cache.get_search("q", "h") is None
assert await cache.get_extract("url") is None
@pytest.mark.asyncio
async def test_disabled_batch_returns_none(self):
config = AppConfig()
config.cache.enabled = False
cache = CacheService(config)
result = await cache.get_extract_batch(["u1", "u2"])
assert result == {"u1": None, "u2": None}
@pytest.mark.asyncio
async def test_disabled_set_is_noop(self):
config = AppConfig()
config.cache.enabled = False
cache = CacheService(config)
# Should not raise
await cache.set_search("q", "h", {"data": 1})
await cache.set_extract("url", "content")
await cache.set_extract_batch([("u1", "c1")])
def test_hash_key_deterministic(self):
key1 = CacheService._hash_key("prefix", "value")
key2 = CacheService._hash_key("prefix", "value")
assert key1 == key2
assert key1.startswith("prefix:")
def test_hash_key_different_values(self):
key1 = CacheService._hash_key("p", "a")
key2 = CacheService._hash_key("p", "b")
assert key1 != key2
# ========== Extractor tests ==========
class TestContentExtractor:
def _make_extractor(self, **kwargs) -> ContentExtractor:
config = AppConfig(**kwargs) if kwargs else AppConfig()
client = MagicMock(spec=httpx.AsyncClient)
return ContentExtractor(config, client)
def test_domain_semaphore_lru_eviction(self):
config = AppConfig()
config.extraction.domain_semaphore_max_size = 3
config.extraction.domain_concurrency = 2
ext = ContentExtractor(config, MagicMock())
# Fill up
ext._get_domain_semaphore("https://a.com/1")
ext._get_domain_semaphore("https://b.com/1")
ext._get_domain_semaphore("https://c.com/1")
assert len(ext._domain_semaphores) == 3
# Adding a 4th should evict the oldest (a.com)
ext._get_domain_semaphore("https://d.com/1")
assert len(ext._domain_semaphores) == 3
assert "a.com" not in ext._domain_semaphores
assert "d.com" in ext._domain_semaphores
def test_domain_semaphore_lru_access_refreshes(self):
config = AppConfig()
config.extraction.domain_semaphore_max_size = 3
ext = ContentExtractor(config, MagicMock())
ext._get_domain_semaphore("https://a.com/1")
ext._get_domain_semaphore("https://b.com/1")
ext._get_domain_semaphore("https://c.com/1")
# Access a.com again to refresh it
ext._get_domain_semaphore("https://a.com/2")
# Now adding d.com should evict b.com (oldest non-refreshed)
ext._get_domain_semaphore("https://d.com/1")
assert "a.com" in ext._domain_semaphores
assert "b.com" not in ext._domain_semaphores
def test_get_headers_returns_valid_headers(self):
ext = self._make_extractor()
headers = ext._get_headers()
assert "User-Agent" in headers
assert "Accept" in headers
assert "Mozilla" in headers["User-Agent"]
def test_to_markdown_with_title(self):
html = "<html><head><title>Test Page</title></head><body></body></html>"
result = ContentExtractor._to_markdown("Some text", html, "https://example.com")
assert "# Test Page" in result
assert "Source: https://example.com" in result
assert "---" in result
assert "Some text" in result
def test_to_markdown_without_title(self):
html = "<html><head></head><body></body></html>"
result = ContentExtractor._to_markdown("Text only", html, "https://example.com")
assert "# " not in result
assert "Source: https://example.com" in result
assert "Text only" in result
# ========== Search backend tests ==========
class TestFallbackSearchBackend:
@pytest.mark.asyncio
async def test_fallback_on_primary_failure(self):
primary = MagicMock(spec=SearXNGBackend)
primary.search = AsyncMock(side_effect=RuntimeError("SearXNG down"))
fallback = MagicMock(spec=DuckDuckGoBackend)
fallback_resp = BackendSearchResponse(
results=[RawSearchResult(title="DDG", url="https://ddg.com", snippet="fallback", score=0.5)],
)
fallback.search = AsyncMock(return_value=fallback_resp)
backend = FallbackSearchBackend(primary, fallback)
result = await backend.search("test query", max_results=1)
primary.search.assert_called_once()
fallback.search.assert_called_once()
assert result.results[0].title == "DDG"
@pytest.mark.asyncio
async def test_no_fallback_when_primary_succeeds(self):
primary_resp = BackendSearchResponse(
results=[RawSearchResult(title="Primary", url="https://p.com", snippet="ok", score=0.9)],
)
primary = MagicMock()
primary.search = AsyncMock(return_value=primary_resp)
fallback = MagicMock()
fallback.search = AsyncMock()
backend = FallbackSearchBackend(primary, fallback)
result = await backend.search("test", max_results=1)
primary.search.assert_called_once()
fallback.search.assert_not_called()
assert result.results[0].title == "Primary"
# ========== Reranker tests ==========
class TestRerankerService:
def test_disabled_reranker_passthrough(self):
reranker = RerankerService()
# Not initialized (disabled)
results = [{"title": "A", "content": "a", "score": 0.5}]
output = reranker.rerank("query", results)
assert output == results # unchanged
def test_empty_results_passthrough(self):
reranker = RerankerService()
output = reranker.rerank("query", [])
assert output == []

97
tests/test_stream.py Normal file
View file

@ -0,0 +1,97 @@
from __future__ import annotations
import json
import pytest
@pytest.mark.asyncio
async def test_stream_basic(client):
resp = await client.post(
"/search/stream",
json={"query": "test stream", "max_results": 2},
)
assert resp.status_code == 200
assert "text/event-stream" in resp.headers.get("content-type", "")
# Parse SSE events
events = _parse_sse(resp.text)
result_events = [e for e in events if e["event"] == "result"]
done_events = [e for e in events if e["event"] == "done"]
assert len(result_events) == 2
assert len(done_events) == 1
# Verify result structure
first = json.loads(result_events[0]["data"])
assert "title" in first
assert "url" in first
assert "content" in first
assert "score" in first
# Verify done event
done_data = json.loads(done_events[0]["data"])
assert "response_time" in done_data
@pytest.mark.asyncio
async def test_stream_with_images(client):
resp = await client.post(
"/search/stream",
json={"query": "cats", "include_images": True, "max_results": 1},
)
assert resp.status_code == 200
events = _parse_sse(resp.text)
image_events = [e for e in events if e["event"] == "image"]
assert len(image_events) > 0
img = json.loads(image_events[0]["data"])
assert "url" in img
assert "description" in img
@pytest.mark.asyncio
async def test_stream_advanced_with_extraction(client):
resp = await client.post(
"/search/stream",
json={"query": "test", "search_depth": "advanced", "max_results": 2},
)
assert resp.status_code == 200
events = _parse_sse(resp.text)
extraction_events = [e for e in events if e["event"] == "extraction"]
assert len(extraction_events) > 0
ext = json.loads(extraction_events[0]["data"])
assert "url" in ext
assert "raw_content" in ext
@pytest.mark.asyncio
async def test_stream_backend_error(client_failing_backend):
resp = await client_failing_backend.post(
"/search/stream",
json={"query": "test", "max_results": 1},
)
assert resp.status_code == 200 # SSE always returns 200
events = _parse_sse(resp.text)
error_events = [e for e in events if e["event"] == "error"]
assert len(error_events) == 1
def _parse_sse(text: str) -> list[dict[str, str]]:
"""Parse Server-Sent Events text into a list of {event, data} dicts."""
events: list[dict[str, str]] = []
current: dict[str, str] = {}
for line in text.split("\n"):
line = line.strip()
if line.startswith("event:"):
current["event"] = line[len("event:"):].strip()
elif line.startswith("data:"):
current["data"] = line[len("data:"):].strip()
elif line == "" and current:
if "event" in current and "data" in current:
events.append(current)
current = {}
if "event" in current and "data" in current:
events.append(current)
return events