mirror of
https://github.com/MaxRyabov/orio-search.git
synced 2026-10-19 15:23:33 +00:00
init
This commit is contained in:
commit
eeaaa9252e
37 changed files with 3839 additions and 0 deletions
140
.gitignore
vendored
Normal file
140
.gitignore
vendored
Normal 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
31
Dockerfile
Normal 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
422
README.md
Normal 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
0
app/__init__.py
Normal file
25
app/auth.py
Normal file
25
app/auth.py
Normal 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
115
app/config.py
Normal 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
48
app/logging_setup.py
Normal 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
200
app/main.py
Normal 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
36
app/middleware.py
Normal 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
0
app/models/__init__.py
Normal file
97
app/models/schemas.py
Normal file
97
app/models/schemas.py
Normal 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
21
app/rate_limit.py
Normal 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
0
app/routers/__init__.py
Normal file
75
app/routers/extract.py
Normal file
75
app/routers/extract.py
Normal 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
122
app/routers/search.py
Normal 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
|
||||
81
app/routers/search_stream.py
Normal file
81
app/routers/search_stream.py
Normal 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
0
app/services/__init__.py
Normal file
121
app/services/cache.py
Normal file
121
app/services/cache.py
Normal 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
172
app/services/extractor.py
Normal 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
52
app/services/reranker.py
Normal 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
|
||||
53
app/services/resilience.py
Normal file
53
app/services/resilience.py
Normal 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]
|
||||
257
app/services/search_backend.py
Normal file
257
app/services/search_backend.py
Normal 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
70
config.yaml
Normal 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
50
docker-compose.yml
Normal 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
387
quick_test.http
Normal 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
20
requirements.txt
Normal 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
13
searxng/settings.yml
Normal 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
0
tests/__init__.py
Normal file
127
tests/conftest.py
Normal file
127
tests/conftest.py
Normal 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
141
tests/test_auth.py
Normal 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
103
tests/test_config.py
Normal 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
153
tests/test_extract.py
Normal 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
62
tests/test_health.py
Normal 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
151
tests/test_schemas.py
Normal 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
214
tests/test_search.py
Normal 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
183
tests/test_services.py
Normal 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
97
tests/test_stream.py
Normal 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
|
||||
Loading…
Add table
Add a link
Reference in a new issue