clossers/didi-lot1-ai
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Livrare LOT 1 - Didi

Browse source
This commit is contained in:
Dezvoltari Evotech 2026-06-25 14:13:25 -07:00
commit 5380c3fc63
990 changed files with 133308 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

123
ai_platform/modules/rerank/.env.example Normal file
View file

@ -0,0 +1,123 @@
# Rerank Configuration
# Copy this file to .env and fill in ALL required values
# The application will fail to start if required variables are missing
# =============================================================================
# REQUIRED: Backend Selection (no defaults - must be set explicitly)
# =============================================================================
# Default backend: vllm or llamacpp
RERANK_DEFAULT_BACKEND=vllm
# Enable backends (true/false)
RERANK_ENABLE_VLLM=true
RERANK_ENABLE_LLAMACPP=false
# External URL for OpenAPI spec (REQUIRED)
RERANK_EXTERNAL_URL=http://localhost:14200
# =============================================================================
# OPTIONAL: API Server Settings
# =============================================================================
# Server binding
# RERANK_HOST=0.0.0.0
# RERANK_PORT=14200 # Prod: 14200, Dev: 54200
# =============================================================================
# OPTIONAL: API Authentication
# =============================================================================
# API tokens for Bearer authentication (comma-separated)
# If not set, authentication is disabled and all endpoints are public
# If set, requests to protected endpoints require: Authorization: Bearer <token>
# Health endpoints (/health, /ready) are always public
# RERANK_API_TOKENS=token1,token2,token3
# =============================================================================
# OPTIONAL: vLLM Backend Configuration
# =============================================================================
# vLLM server URL (internal Docker network)
# RERANK_VLLM_BASE_URL=http://localhost:54201
# vLLM model to load (Hugging Face model ID for cross-encoder)
RERANK_VLLM_MODEL=BAAI/bge-reranker-v2-m3
# vLLM server port
# RERANK_VLLM_PORT=14201
# GPU assignment for vLLM
# RERANK_VLLM_GPU=0
# GPU memory utilization
# RERANK_VLLM_GPU_UTIL=0.50
# Maximum model sequence length
# RERANK_VLLM_MAX_LEN=8192
# =============================================================================
# OPTIONAL: llama.cpp Backend Configuration
# =============================================================================
# llama.cpp server URL (internal Docker network)
# RERANK_LLAMACPP_BASE_URL=http://localhost:54210
# Directory containing model files (absolute path recommended)
MODELS_DIR=/path/to/models
# llama.cpp model file (GGUF format, filename only - must be in MODELS_DIR)
RERANK_LLAMACPP_MODEL=bge-reranker-v2-m3-q4_k_m.gguf
# llama.cpp server port
# RERANK_LLAMACPP_PORT=14210
# llama.cpp context size
# RERANK_LLAMACPP_CTX=8192
# llama.cpp performance settings
# RERANK_LLAMACPP_THREADS=4
# RERANK_LLAMACPP_PARALLEL=4
# =============================================================================
# OPTIONAL: Timeout Settings
# =============================================================================
# Request timeout for rerank calls (seconds)
# RERANK_REQUEST_TIMEOUT=120.0
# Connection timeout (seconds)
# RERANK_CONNECT_TIMEOUT=10.0
# =============================================================================
# OPTIONAL: Rate Limiting and Concurrency
# =============================================================================
# Requests per second limit
# RERANK_RATE_LIMIT_RPS=20.0
# Maximum burst size for rate limiting
# RERANK_RATE_LIMIT_BURST=50
# Maximum concurrent rerank requests
# RERANK_MAX_CONCURRENT_RERANKS=20
# =============================================================================
# OPTIONAL: Logging
# =============================================================================
# Log level: DEBUG, INFO, WARNING, ERROR, CRITICAL
# RERANK_LOG_LEVEL=INFO
# Enable JSON logging format (true/false)
# RERANK_LOG_JSON=false
# =============================================================================
# OPTIONAL: Shared Settings (for Docker)
# =============================================================================
# HuggingFace cache directory
# HF_CACHE_DIR=/cai2_ds_storage/hf_cache
# HuggingFace token
# HF_TOKEN=

376
ai_platform/modules/rerank/API.md Normal file
View file

@ -0,0 +1,376 @@
# Rerank API Reference
Cohere/Jina-compatible reranking API with multiple backend support.
## Base URL
```
{BASE_URL}
```
- **Local development:** `http://localhost:54200`
- **Docker (internal):** `http://didiAI-rerank-api:14200`
- **Production:** Use your configured hostname
## Authentication
Authentication is **optional**. If `RERANK_API_TOKENS` is set, requests require a Bearer token:
```
Authorization: Bearer <token>
```
Health endpoints (`/health`, `/ready`) are always public.
## Endpoints
### Rerank Documents
Rerank documents against a query based on relevance.
**Endpoints:**
- `POST /v1/rerank`
- `POST /v2/rerank` (alias)
**Request Headers:**
| Header | Required | Description |
|--------|----------|-------------|
| `Content-Type` | Yes | Must be `application/json` |
| `Authorization` | If auth enabled | `Bearer <token>` |
**Request Body:**
```json
{
"model": "BAAI/bge-reranker-v2-m3",
"query": "What is machine learning?",
"documents": [
"Machine learning is a subset of AI",
"Cats are pets",
"Deep learning uses neural networks"
],
"top_n": 3,
"return_documents": false,
"backend": null
}
```
| Field | Type | Required | Description |
|-------|------|----------|-------------|
| `model` | string | Yes | Model identifier for reranking |
| `query` | string | Yes | The search query |
| `documents` | string[] | Yes | Documents to rerank (1-1000) |
| `top_n` | integer | No | Number of results to return (default: all) |
| `return_documents` | boolean | No | Include document text in response |
| `backend` | string | No | Override default backend: `"vllm"` or `"llamacpp"` |
**Response:**
```json
{
"id": "rerank-abc123def456",
"model": "BAAI/bge-reranker-v2-m3",
"results": [
{
"index": 2,
"relevance_score": 0.9523,
"document": null
},
{
"index": 0,
"relevance_score": 0.8876,
"document": null
},
{
"index": 1,
"relevance_score": 0.0234,
"document": null
}
],
"usage": {
"total_tokens": 150
},
"backend": "vllm"
}
```
**Example:**
```bash
curl -X POST http://localhost:54200/v1/rerank \
-H "Content-Type: application/json" \
-d '{
"model": "BAAI/bge-reranker-v2-m3",
"query": "What is machine learning?",
"documents": [
"Machine learning is a subset of AI",
"Cats are pets",
"Deep learning uses neural networks"
]
}'
```
**Get top 2 results:**
```bash
curl -X POST http://localhost:54200/v1/rerank \
-H "Content-Type: application/json" \
-d '{
"model": "BAAI/bge-reranker-v2-m3",
"query": "programming languages",
"documents": ["Python", "Java", "Cooking", "C++", "Hiking"],
"top_n": 2
}'
```
**Include documents in response:**
```bash
curl -X POST http://localhost:54200/v1/rerank \
-H "Content-Type: application/json" \
-d '{
"model": "BAAI/bge-reranker-v2-m3",
"query": "AI",
"documents": ["Machine learning", "Deep learning", "Recipes"],
"return_documents": true
}'
```
---
### List Models
List available reranking models.
**Endpoint:** `GET /v1/models`
**Query Parameters:**
| Parameter | Type | Required | Description |
|-----------|------|----------|-------------|
| `backend` | string | No | Filter by backend: `"vllm"` or `"llamacpp"` |
**Response:**
```json
{
"object": "list",
"data": [
{
"id": "BAAI/bge-reranker-v2-m3",
"backend": "vllm",
"loaded": true,
"max_input_tokens": null
}
]
}
```
**Example:**
```bash
# List all models
curl http://localhost:54200/v1/models
# List models from specific backend
curl "http://localhost:54200/v1/models?backend=vllm"
```
---
### List Backends
List available backends.
**Endpoint:** `GET /v1/backends`
**Response:**
```json
{
"backends": ["vllm", "llamacpp"]
}
```
**Example:**
```bash
curl http://localhost:54200/v1/backends
```
---
### Health Check
Detailed health status including per-backend health.
**Endpoint:** `GET /health`
**Response:**
```json
{
"status": "healthy",
"backends": [
{
"name": "vllm",
"healthy": true,
"message": null
}
]
}
```
Status values:
- `"healthy"` - All backends are healthy
- `"degraded"` - Some backends are unhealthy
- `"unhealthy"` - All backends are unhealthy
**Example:**
```bash
curl http://localhost:54200/health
```
---
### Readiness Probe
Simple readiness check for Kubernetes.
**Endpoint:** `GET /ready`
**Response:**
```json
{
"ready": true
}
```
**Example:**
```bash
curl http://localhost:54200/ready
```
---
## Error Responses
All errors follow this format:
```json
{
"detail": "Error message describing what went wrong"
}
```
### HTTP Status Codes
| Code | Meaning |
|------|---------|
| 200 | Success |
| 400 | Bad request (invalid parameters, backend not enabled) |
| 401 | Authentication required or failed |
| 429 | Rate limit exceeded |
| 503 | Service unavailable (backend connection failed) |
| 504 | Gateway timeout (backend request timed out) |
### Rate Limit Response
```json
{
"detail": "Too many requests",
"retry_after": 1.5
}
```
Headers include: `Retry-After: 2`
---
## Request Headers
| Header | Required | Description |
|--------|----------|-------------|
| `Content-Type` | Yes (POST) | Must be `application/json` |
| `Authorization` | If auth enabled | `Bearer <token>` |
| `X-Request-ID` | No | Request tracking ID (generated if not provided) |
Response always includes `X-Request-ID` header for tracking.
---
## SDK Examples
### Python (httpx)
```python
import httpx
async def rerank_documents(
query: str,
documents: list[str],
top_n: int | None = None,
) -> list[tuple[int, float]]:
async with httpx.AsyncClient() as client:
response = await client.post(
"http://localhost:54200/v1/rerank",
json={
"model": "BAAI/bge-reranker-v2-m3",
"query": query,
"documents": documents,
"top_n": top_n,
},
headers={"Authorization": "Bearer your-token"},
)
response.raise_for_status()
data = response.json()
return [(r["index"], r["relevance_score"]) for r in data["results"]]
```
### Python (requests)
```python
import requests
def rerank_documents(query: str, documents: list[str]) -> dict:
response = requests.post(
"http://localhost:54200/v1/rerank",
json={
"model": "BAAI/bge-reranker-v2-m3",
"query": query,
"documents": documents,
},
)
response.raise_for_status()
return response.json()
# Usage
result = rerank_documents(
"What is AI?",
["Machine learning", "Deep learning", "Cooking recipes"]
)
for r in result["results"]:
print(f"Doc {r['index']}: {r['relevance_score']:.4f}")
```
### curl
```bash
# Basic reranking
curl -X POST http://localhost:54200/v1/rerank \
-H "Content-Type: application/json" \
-H "Authorization: Bearer your-token" \
-d '{
"model": "BAAI/bge-reranker-v2-m3",
"query": "programming",
"documents": ["Python", "Java", "Cooking"]
}'
# With specific backend
curl -X POST http://localhost:54200/v1/rerank \
-H "Content-Type: application/json" \
-d '{
"model": "BAAI/bge-reranker-v2-m3",
"query": "programming",
"documents": ["Python", "Java"],
"backend": "vllm"
}'
# v2 endpoint (identical)
curl -X POST http://localhost:54200/v2/rerank \
-H "Content-Type: application/json" \
-d '{
"model": "BAAI/bge-reranker-v2-m3",
"query": "test",
"documents": ["doc1", "doc2"]
}'
```

274
ai_platform/modules/rerank/INDEX.md Normal file
View file

@ -0,0 +1,274 @@
# Rerank Service — INDEX
Cross-encoder reranking API for DIDI semantic search precision. Serves the `BAAI/bge-reranker-v2-m3` model behind a Cohere/Jina-compatible HTTP API. Used after kNN retrieval to score `(query, candidate)` pairs and reorder evidence by true relevance, instead of relying solely on pgvector cosine similarity from embeddings.
- **Stack:** Python 3.10+, FastAPI, Uvicorn, Pydantic v2, httpx, vLLM (or llama.cpp) backend
- **URLs:**
- Dev API: `http://10.11.10.12:54200`
- Prod API: `http://10.11.10.12:14200`
- vLLM internal: `:54201` (Dev) / `:14201` (Prod)
- llama.cpp internal: `:54210` (Dev) / `:14210` (Prod)
- **Containers:** `didiAI-rerank-api` (FastAPI wrapper) + `didiAI-rerank-vllm` (or `didiAI-rerank-llamacpp`)
- **Model:** `BAAI/bge-reranker-v2-m3` (multilingual cross-encoder, ~8K context window)
- **Port schema:** `x42xx` family (Reranking)
---
## Ce face
Re-scores candidate documents against a query using a cross-encoder model — directly attending over both texts together — for higher precision than embedding cosine similarity alone. The cross-encoder gives a true relevance score in `[0..1]`.
**Pipeline role inside DIDI:**
1. didi-brain `/v1/gather` runs an initial kNN retrieval over pgvector embeddings (dim 1024, separate `embeddings` module).
2. The top-K candidates (typically 50–200) are passed to this rerank service.
3. The reranker scores each `(query, candidate)` pair and returns a relevance-sorted list.
4. didi-brain takes the top-N (e.g. top 10) reranked items as final evidence.
This two-stage retrieval (kNN → cross-encoder rerank) is the standard high-quality semantic search pattern: embeddings handle scale, the cross-encoder handles precision.
---
## API endpoints
Cohere/Jina-compatible. Auth is optional (Bearer token) — enabled if `RERANK_API_TOKENS` is set. Health endpoints stay public.
### `POST /v1/rerank` (alias `POST /v2/rerank`)
Rerank a list of documents against a query.
**Request body:**
```json
{
"model": "BAAI/bge-reranker-v2-m3",
"query": "What is machine learning?",
"documents": ["Machine learning is...", "Cats are pets", "Deep learning..."],
"top_n": 3,
"return_documents": false,
"backend": null
}
```
| Field | Type | Required | Description |
|-------|------|----------|-------------|
| `model` | string | yes | Model identifier (e.g. `BAAI/bge-reranker-v2-m3`) |
| `query` | string | yes | Search query |
| `documents` | string[] | yes | Documents to rerank (1–1000) |
| `top_n` | int | no | Return only top N results (default: all) |
| `return_documents` | bool | no | Echo document text in response |
| `backend` | string | no | Override default: `"vllm"` or `"llamacpp"` |
**Response:**
```json
{
"id": "rerank-abc123def456",
"model": "BAAI/bge-reranker-v2-m3",
"results": [
{"index": 2, "relevance_score": 0.9523, "document": null},
{"index": 0, "relevance_score": 0.8876, "document": null}
],
"usage": {"total_tokens": 150},
"backend": "vllm"
}
```
### `GET /v1/models`
List loaded models. Optional `?backend=vllm|llamacpp` filter.
### `GET /v1/backends`
List enabled backends, e.g. `{"backends": ["vllm", "llamacpp"]}`.
### `GET /health`
Detailed per-backend health. Status is `healthy` / `degraded` / `unhealthy`.
### `GET /ready`
Simple K8s-style readiness probe — returns `{"ready": true}`.
### Error codes
`400` bad request · `401` auth failed · `429` rate-limited (with `Retry-After` header) · `503` backend unavailable · `504` backend timeout. All errors return `{"detail": "..."}`. Every response carries an `X-Request-ID` header for tracing.
Full reference in `API.md`.
---
## Backends
The API is a thin FastAPI router in front of one of two cross-encoder servers. Both backends speak OpenAI-compatible HTTP, so the wrapper unifies them.
| Backend | When | Pros | Cons |
|---------|------|------|------|
| **vLLM** | Production, GPU host | High throughput, batched scoring, lowest latency under load | Requires NVIDIA GPU + CUDA 12.x |
| **llama.cpp** | Dev / CPU fallback | Runs on CPU or modest GPU, GGUF quantized models, low memory | Lower throughput |
The wrapper routes requests via a backend registry. Set `RERANK_DEFAULT_BACKEND` to choose, or override per-request with the `backend` field in the body. Both can be enabled simultaneously (`RERANK_ENABLE_VLLM=true`, `RERANK_ENABLE_LLAMACPP=true`).
vLLM is launched with `--task score` (vLLM 0.8.x cross-encoder mode). llama.cpp is launched with `--reranking`.
---
## How didi-brain uses it
didi-brain's gather/retrieval flow:
1. `services/gather.py` builds an initial candidate set via pgvector kNN over the `embeddings` module (1024-dim vectors).
2. Calls the reranker via `shared/reranker_client.py` — typically wrapping the `/v1/rerank` endpoint with the brain's HTTPS/auth config.
3. Picks the top-N reranked candidates as final evidence chunks.
4. Passes them to the LLM as grounded context.
The reranker is therefore in the critical path of every brain `/v1/gather` call. Its latency budget is small (a few hundred ms), so vLLM batching matters in production.
---
## Structura fisiere
```
rerank/
├── README.md # Quick start, install, env vars table
├── API.md # Full HTTP API reference + curl/Python examples
├── INDEX.md # This file
├── pyproject.toml # Package metadata, deps (fastapi, httpx, openai SDK)
├── uv.lock # uv lockfile
├── .env.example # All RERANK_* env vars documented
├── deploy/
│ ├── Dockerfile # API wrapper image
│ ├── docker-compose.yml # api / vllm / llamacpp profiles
│ └── deploy.sh # Helper: ./deploy.sh --profile vllm -d
├── src/rerank/
│ ├── __init__.py # Public exports (RerankClient, ...)
│ ├── cli.py # `rerank` entrypoint — `python -m rerank.cli`
│ ├── client.py # RerankClient (Python SDK to call this API)
│ ├── config.py # Pydantic Settings, RERANK_ env prefix
│ ├── schemas.py # Pydantic request/response models
│ ├── exceptions.py # Custom error types
│ ├── logging.py # Structured logging setup
│ ├── types.py # Backend literal types
│ ├── api/
│ │ ├── app.py # FastAPI app factory
│ │ ├── dependencies.py # Auth + rate-limit DI
│ │ ├── middleware.py # X-Request-ID, rate limit, error handlers
│ │ └── routes/
│ │ ├── rerank.py # POST /v1/rerank, /v2/rerank
│ │ ├── models.py # GET /v1/models, /v1/backends
│ │ └── health.py # GET /health, /ready
│ └── backends/
│ ├── base.py # Abstract BackendBase (rerank, health, list_models)
│ ├── vllm_backend.py # vLLM via OpenAI SDK (--task score)
│ ├── llamacpp_backend.py # llama.cpp --reranking endpoint
│ └── registry.py # Backend registry + selector
└── tests/
├── conftest.py
├── test_config.py
└── test_schemas.py
```
---
## Configuration
All env vars use the `RERANK_` prefix and are read via Pydantic Settings (`config.py`).
**Required (no defaults):**
| Var | Example |
|-----|---------|
| `RERANK_DEFAULT_BACKEND` | `vllm` or `llamacpp` |
| `RERANK_ENABLE_VLLM` | `true` / `false` |
| `RERANK_ENABLE_LLAMACPP` | `true` / `false` |
| `RERANK_EXTERNAL_URL` | `http://localhost:14200` (used in OpenAPI spec) |
**API server:**
| Var | Default | Notes |
|-----|---------|-------|
| `RERANK_HOST` | `0.0.0.0` | |
| `RERANK_PORT` | `14200` | Prod 14200 / Dev 54200 |
| `RERANK_API_TOKENS` | (empty) | Comma-separated; if empty, auth is disabled |
| `RERANK_RATE_LIMIT_RPS` | `20.0` | |
| `RERANK_RATE_LIMIT_BURST` | `50` | |
| `RERANK_MAX_CONCURRENT_RERANKS` | `20` | |
| `RERANK_REQUEST_TIMEOUT` | `120.0` s | |
| `RERANK_CONNECT_TIMEOUT` | `10.0` s | |
| `RERANK_LOG_LEVEL` | `INFO` | |
| `RERANK_LOG_JSON` | `false` | |
**vLLM backend:**
| Var | Default |
|-----|---------|
| `RERANK_VLLM_BASE_URL` | `http://localhost:54201` |
| `RERANK_VLLM_MODEL` | `BAAI/bge-reranker-v2-m3` |
| `RERANK_VLLM_PORT` | `14201` |
| `RERANK_VLLM_GPU` | `0` |
| `RERANK_VLLM_GPU_UTIL` | `0.50` |
| `RERANK_VLLM_MAX_LEN` | `8192` |
**llama.cpp backend:**
| Var | Default |
|-----|---------|
| `RERANK_LLAMACPP_BASE_URL` | `http://localhost:54210` |
| `RERANK_LLAMACPP_MODEL` | `bge-reranker-v2-m3-q4_k_m.gguf` |
| `RERANK_LLAMACPP_PORT` | `14210` |
| `RERANK_LLAMACPP_CTX` | `8192` |
| `RERANK_LLAMACPP_THREADS` | `4` |
| `RERANK_LLAMACPP_PARALLEL` | `4` |
| `MODELS_DIR` | `/cai2_ds_storage/models` |
**Shared HF:** `HF_CACHE_DIR` (default `/cai2_ds_storage/hf_cache`), `HF_TOKEN`.
Full list with comments in `.env.example`.
---
## Performance / model
- **Model:** `BAAI/bge-reranker-v2-m3` — multilingual cross-encoder, ~8K context window, ~568M params.
- **Output:** single relevance score per `(query, document)` pair, normalized to `[0..1]`.
- **Expected latency:** roughly 100–300 ms for ~20 candidates per query on a single mid-range GPU with vLLM batching; scales sublinearly with batch size.
- **Throughput:** dominated by GPU memory and `RERANK_VLLM_GPU_UTIL` / `RERANK_MAX_CONCURRENT_RERANKS`. vLLM continuous batching helps a lot under concurrent load.
- **Memory:** vLLM defaults to 50% GPU memory util (`RERANK_VLLM_GPU_UTIL=0.50`), so it can co-host on a GPU shared with other modules. llama.cpp Q4_K_M quantization fits comfortably on CPU.
---
## Deployment
Docker Compose with three profiles in `deploy/docker-compose.yml`:
```bash
cd /home/admin365/didi_mono/ai_platform/modules/rerank/deploy
# Production (GPU + vLLM)
./deploy.sh --profile vllm -d
# Lightweight / no-GPU
./deploy.sh --profile llamacpp -d
# API only (external rerank servers running elsewhere)
./deploy.sh --profile api -d
# Logs / stop
./deploy.sh --profile vllm --logs
./deploy.sh --profile vllm --down
```
Containers:
- `didiAI-rerank-api` — FastAPI wrapper, port `RERANK_PORT` (14200/54200)
- `didiAI-rerank-vllm` — `vllm/vllm-openai:v0.8.5`, GPU device pinned via `RERANK_VLLM_GPU`, internal port 14201
- `didiAI-rerank-llamacpp` — `ghcr.io/ggml-org/llama.cpp:server-b4769`, internal port 8080 → host 14210/54210
All services share the external `didi-network` Docker network so other DIDI modules (didi-brain, embeddings) can reach them by container name.
Health gates: API has a 30-s `/health` healthcheck; vLLM has a 5-min start period to allow model load on first boot.
---
## Related
- **Embeddings module** (`ai_platform/modules/embeddings`) — separate service, dim 1024, used for the initial pgvector kNN stage that feeds this reranker.
- **didi-brain** (`ai_platform/modules/didi_brain/brain_api`) — main consumer; `services/gather.py` calls rerank after kNN; `shared/reranker_client.py` is the helper.
- **GPU host:** `10.11.10.17` / `10.11.10.12` (see project network architecture doc).

206
ai_platform/modules/rerank/README.md Normal file
View file

@ -0,0 +1,206 @@
# Rerank Module
Cohere/Jina-compatible reranking API with support for vLLM and llama.cpp backends.
## Features
- **Cohere/Jina-compatible API**: Drop-in replacement for Cohere's `/v1/rerank` endpoint
- **Multiple backends**: Support for vLLM and llama.cpp
- **High performance**: Built on FastAPI with async support
- **Production-ready**: Rate limiting, authentication, health checks
## Prerequisites
### Required
- **Linux** - Ubuntu 22.04+ or similar
- **Python 3.10+** - Managed via `uv`
- **uv** - Fast Python package manager
### Optional (for backends)
- **vLLM** - Requires NVIDIA GPU with CUDA 12.x
- **llama.cpp** - Can run on CPU or GPU
## Quick Start
### 1. Install dependencies
```bash
cd modules/rerank
uv sync --all-extras
```
### 2. Configure environment
```bash
cp .env.example .env
# Edit .env with your settings
```
### 3. Start backend server
**vLLM (GPU):**
```bash
python -m vllm.entrypoints.openai.api_server \
--model BAAI/bge-reranker-v2-m3 \
--host 0.0.0.0 --port 54201 \
--task score
```
**llama.cpp (CPU):**
```bash
llama-server \
--model /models/bge-reranker-v2-m3-q4_k_m.gguf \
--host 0.0.0.0 --port 54210 \
--reranking
```
### 4. Start API server
```bash
# Set required environment variables
export RERANK_DEFAULT_BACKEND=vllm
export RERANK_ENABLE_VLLM=true
export RERANK_ENABLE_LLAMACPP=false
export RERANK_EXTERNAL_URL=http://localhost:54200
# Run the server
uv run python -m rerank.cli --port 54200
```
### 5. Test the API
```bash
curl -X POST http://localhost:54200/v1/rerank \
-H "Content-Type: application/json" \
-d '{
"model": "BAAI/bge-reranker-v2-m3",
"query": "What is machine learning?",
"documents": [
"Machine learning is a subset of AI",
"Cats are pets",
"Deep learning uses neural networks"
]
}'
```
## Docker Deployment
```bash
cd modules/rerank/deploy
# Copy and configure .env
cp ../.env.example .env
# Edit .env with your settings
# Start with vLLM backend
./deploy.sh --profile vllm -d
# Or start with llama.cpp backend
./deploy.sh --profile llamacpp -d
# View logs
./deploy.sh --profile vllm --logs
# Stop
./deploy.sh --profile vllm --down
```
## Python Library Usage
```python
from rerank import RerankClient
# Initialize client (reads config from environment)
client = RerankClient()
# Rerank documents
response = await client.rerank(
query="What is machine learning?",
documents=[
"ML is a subset of AI",
"Cats are pets",
"Deep learning uses neural networks",
],
model="BAAI/bge-reranker-v2-m3",
)
# Access results
for result in response.results:
print(f"Index {result.index}: {result.relevance_score:.4f}")
# Get top N results
response = await client.rerank(
query="programming",
documents=["Python", "Java", "Cooking", "C++"],
model="BAAI/bge-reranker-v2-m3",
top_n=2,
)
# Include documents in response
response = await client.rerank(
query="AI",
documents=["Machine learning", "Deep learning"],
model="BAAI/bge-reranker-v2-m3",
return_documents=True,
)
for result in response.results:
print(f"{result.document}: {result.relevance_score:.4f}")
```
## Configuration
All configuration is via environment variables with the `RERANK_` prefix:
| Variable | Required | Default | Description |
|----------|----------|---------|-------------|
| `RERANK_DEFAULT_BACKEND` | Yes | - | Default backend: `vllm` or `llamacpp` |
| `RERANK_ENABLE_VLLM` | Yes | - | Enable vLLM backend |
| `RERANK_ENABLE_LLAMACPP` | Yes | - | Enable llama.cpp backend |
| `RERANK_EXTERNAL_URL` | Yes | - | External URL for OpenAPI spec |
| `RERANK_PORT` | No | 54200 | API server port |
| `RERANK_VLLM_BASE_URL` | No | http://localhost:54201 | vLLM server URL |
| `RERANK_LLAMACPP_BASE_URL` | No | http://localhost:54210 | llama.cpp server URL |
| `RERANK_API_TOKENS` | No | - | Comma-separated API tokens |
| `RERANK_RATE_LIMIT_RPS` | No | 20.0 | Requests per second limit |
| `RERANK_MAX_CONCURRENT_RERANKS` | No | 20 | Max concurrent requests |
See `.env.example` for the complete list.
## Port Allocation
Following the datacenter port schema (x42xx = Reranking):
| Port | Service | Environment |
|------|---------|-------------|
| 14200 | Rerank API | Production |
| 54200 | Rerank API | Development |
| 14201 | vLLM Rerank Server | Production |
| 54201 | vLLM Rerank Server | Development |
| 14210 | llama.cpp Rerank Server | Production |
| 54210 | llama.cpp Rerank Server | Development |
## Development
```bash
# Install dev dependencies
uv sync --all-extras
# Run tests
uv run pytest
# Run tests with coverage
uv run pytest --cov=src/rerank --cov-report=term-missing
# Lint and format
uv run ruff check .
uv run ruff format .
# Type check
uv run mypy src/
```
## API Reference
See [API.md](API.md) for the complete API documentation.

43
ai_platform/modules/rerank/deploy/Dockerfile Normal file
View file

@ -0,0 +1,43 @@
# Rerank API Server
# Multi-stage build for smaller final image
FROM python:3.11.12-slim AS builder
# Install uv
COPY --from=ghcr.io/astral-sh/uv:0.10 /uv /usr/local/bin/uv
WORKDIR /app
# Copy project files
COPY pyproject.toml README.md ./
COPY src/ ./src/
# Install dependencies (with optional extras for backends)
RUN uv sync --frozen --no-dev --all-extras || uv sync --no-dev --all-extras
# Production image
FROM python:3.11.12-slim
WORKDIR /app
# Copy virtual environment from builder
COPY --from=builder /app/.venv /app/.venv
# Copy source code
COPY src/ ./src/
# Set environment variables
ENV PATH="/app/.venv/bin:$PATH"
ENV PYTHONUNBUFFERED=1
ENV RERANK_HOST=0.0.0.0
ENV RERANK_PORT=14200
# Health check (using Python since curl not available in slim image)
HEALTHCHECK --interval=30s --timeout=10s --retries=3 \
CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:14200/health')" || exit 1
# Expose port
EXPOSE 14200
# Run the server
CMD ["python", "-m", "rerank.cli"]

142
ai_platform/modules/rerank/deploy/deploy.sh Normal file
View file

@ -0,0 +1,142 @@
#!/usr/bin/env bash
#
# Docker Compose Startup Script for Rerank
#
# Usage: ./deploy/deploy.sh [OPTIONS]
#
# Options:
# --profile <api|vllm|llamacpp> Docker compose profile
# --detach Run in detached mode
# --down Stop and remove containers
# --logs Show logs
# --help Show this help message
#
# Required: Set environment variables in .env file or export them before running.
# See .env.example for the full list of required variables.
set -euo pipefail
# Script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# Load .env file if it exists (check deploy/ first, then module root)
ENV_FILE=""
if [[ -f "$SCRIPT_DIR/.env" ]]; then
ENV_FILE="$SCRIPT_DIR/.env"
elif [[ -f "$SCRIPT_DIR/../.env" ]]; then
ENV_FILE="$SCRIPT_DIR/../.env"
fi
if [[ -n "$ENV_FILE" ]]; then
echo "Loading environment from: $ENV_FILE"
set -a
source "$ENV_FILE"
set +a
fi
PROFILE=""
DETACH=""
ACTION="up"
show_help() {
sed -n '2,14p' "$0" | sed 's/^# //' | sed 's/^#//'
exit 0
}
check_required_var() {
local var_name="$1"
if [[ -z "${!var_name:-}" ]]; then
echo "ERROR: Required environment variable $var_name is not set"
echo "Set it in .env file or export it before running this script"
exit 1
fi
}
# Parse arguments
while [[ $# -gt 0 ]]; do
case $1 in
--profile)
PROFILE="$2"
shift 2
;;
--detach|-d)
DETACH="-d"
shift
;;
--down)
ACTION="down"
shift
;;
--logs)
ACTION="logs"
shift
;;
--help|-h)
show_help
;;
*)
echo "Unknown option: $1"
echo "Use --help for usage information"
exit 1
;;
esac
done
# Require profile to be specified
if [[ -z "$PROFILE" && "$ACTION" == "up" ]]; then
echo "ERROR: --profile is required"
echo "Options: api, vllm, llamacpp"
exit 1
fi
# Check required variables
check_required_var "RERANK_DEFAULT_BACKEND"
check_required_var "RERANK_ENABLE_VLLM"
check_required_var "RERANK_ENABLE_LLAMACPP"
check_required_var "RERANK_EXTERNAL_URL"
# Check profile-specific variables
case $PROFILE in
vllm)
check_required_var "RERANK_VLLM_MODEL"
;;
llamacpp)
check_required_var "MODELS_DIR"
check_required_var "RERANK_LLAMACPP_MODEL"
;;
esac
cd "$SCRIPT_DIR"
case $ACTION in
up)
echo "Starting Rerank with profile: $PROFILE"
echo " Default backend: $RERANK_DEFAULT_BACKEND"
echo " vLLM enabled: $RERANK_ENABLE_VLLM"
echo " llama.cpp enabled: $RERANK_ENABLE_LLAMACPP"
if [[ "$PROFILE" == "vllm" ]]; then
echo " vLLM model: ${RERANK_VLLM_MODEL:-BAAI/bge-reranker-v2-m3}"
fi
if [[ "$PROFILE" == "llamacpp" ]]; then
echo " llama.cpp model: ${RERANK_LLAMACPP_MODEL:-bge-reranker-v2-m3-q4_k_m.gguf}"
fi
echo ""
# shellcheck disable=SC2086
exec docker compose --profile "$PROFILE" up $DETACH
;;
down)
if [[ -z "$PROFILE" ]]; then
echo "ERROR: --profile is required with --down"
exit 1
fi
echo "Stopping Rerank containers..."
exec docker compose --profile "$PROFILE" down
;;
logs)
if [[ -z "$PROFILE" ]]; then
echo "ERROR: --profile is required with --logs"
exit 1
fi
exec docker compose --profile "$PROFILE" logs -f
;;
esac

131
ai_platform/modules/rerank/deploy/docker-compose.yml Normal file
View file

@ -0,0 +1,131 @@
# Rerank Module - Docker Compose Configuration
#
# Port Allocation (x42xx = Reranking):
# 14200 - Rerank API (Prod)
# 54200 - Rerank API (Dev)
# 14201/54201 - vLLM Rerank Server
# 14210/54210 - llama.cpp Rerank Server
#
# Profiles:
# api - API server only (uses external rerank servers)
# vllm - API + vLLM server (GPU required)
# llamacpp - API + llama.cpp server
#
# Naming Convention: didiAI-{module}-{service}
#
# Network:
# Uses deploy_default network (shared with other modules)
networks:
deploy_default:
external: true
services:
# ==========================================================================
# Rerank API Server
# ==========================================================================
rerank-api:
container_name: didiAI-rerank-api
image: didiai-rerank-api
build:
context: ..
dockerfile: deploy/Dockerfile
ports:
- "${RERANK_PORT:-14200}:${RERANK_PORT:-14200}"
networks:
- deploy_default
environment:
- RERANK_PORT=${RERANK_PORT:-14200}
- RERANK_EXTERNAL_URL=${RERANK_EXTERNAL_URL}
- RERANK_DEFAULT_BACKEND=${RERANK_DEFAULT_BACKEND}
- RERANK_ENABLE_VLLM=${RERANK_ENABLE_VLLM}
- RERANK_ENABLE_LLAMACPP=${RERANK_ENABLE_LLAMACPP}
- RERANK_VLLM_BASE_URL=http://didiAI-rerank-vllm:14201
- RERANK_LLAMACPP_BASE_URL=http://didiAI-rerank-llamacpp:8080
- RERANK_API_TOKENS=${RERANK_API_TOKENS:-}
- RERANK_DASHBOARD_URL=${RERANK_DASHBOARD_URL:-http://didiAI-dashboard:51300}
healthcheck:
test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:${RERANK_PORT:-14200}/health')"]
interval: 30s
timeout: 10s
retries: 3
start_period: 30s
restart: unless-stopped
profiles:
- api
- vllm
- llamacpp
# ==========================================================================
# vLLM Rerank Server
# ==========================================================================
# Cross-encoder model using vLLM with --task score (v0.8.x syntax)
vllm-rerank:
container_name: didiAI-rerank-vllm
image: vllm/vllm-openai:v0.8.5
ports:
- "${RERANK_VLLM_PORT:-14201}:14201"
networks:
- deploy_default
volumes:
- ${HF_CACHE_DIR:-/cai2_ds_storage/hf_cache}:/root/.cache/huggingface
environment:
- HF_HOME=/root/.cache/huggingface
- HUGGING_FACE_HUB_TOKEN=${HF_TOKEN:-}
- CUDA_VISIBLE_DEVICES=${RERANK_VLLM_GPU:-0}
command: >
--model ${RERANK_VLLM_MODEL:-BAAI/bge-reranker-v2-m3}
--host 0.0.0.0
--port 14201
--task score
--trust-remote-code
--max-model-len ${RERANK_VLLM_MAX_LEN:-8192}
--gpu-memory-utilization ${RERANK_VLLM_GPU_UTIL:-0.50}
--disable-log-requests
deploy:
resources:
reservations:
devices:
- driver: nvidia
device_ids: ['${RERANK_VLLM_GPU:-0}']
capabilities: [gpu]
healthcheck:
test: ["CMD", "curl", "-f", "http://localhost:14201/health"]
interval: 30s
timeout: 10s
retries: 10
start_period: 300s
restart: unless-stopped
profiles:
- vllm
# ==========================================================================
# llama.cpp Rerank Server
# ==========================================================================
# Cross-encoder model using llama.cpp with --reranking
llamacpp-rerank:
container_name: didiAI-rerank-llamacpp
image: ghcr.io/ggml-org/llama.cpp:server-b4769
ports:
- "${RERANK_LLAMACPP_PORT:-14210}:8080"
networks:
- deploy_default
volumes:
- ${MODELS_DIR:-/cai2_ds_storage/models}:/models:ro
command: >
--model /models/${RERANK_LLAMACPP_MODEL:-bge-reranker-v2-m3-q4_k_m.gguf}
--host 0.0.0.0
--port 8080
--reranking
--ctx-size ${RERANK_LLAMACPP_CTX:-8192}
--threads ${RERANK_LLAMACPP_THREADS:-4}
--parallel ${RERANK_LLAMACPP_PARALLEL:-4}
healthcheck:
test: ["CMD", "curl", "-f", "http://localhost:8080/health"]
interval: 30s
timeout: 10s
retries: 3
start_period: 60s
restart: unless-stopped
profiles:
- llamacpp

75
ai_platform/modules/rerank/pyproject.toml Normal file
View file

@ -0,0 +1,75 @@
[project]
name = "rerank"
version = "0.1.0"
description = "Cohere/Jina-compatible rerank API with vLLM/llama.cpp backends"
requires-python = ">=3.10"
readme = "README.md"
dependencies = [
# Core dependencies (always installed)
"fastapi>=0.115.0,<1.0",
"uvicorn[standard]>=0.32.0",
"pydantic>=2.0,<3.0",
"pydantic-settings>=2.0,<3.0",
"httpx>=0.27.0,<1.0",
"prometheus-fastapi-instrumentator>=7.0.0",
"opentelemetry-instrumentation-fastapi>=0.50b0",
"opentelemetry-exporter-otlp-proto-grpc>=1.30.0",
]
[project.optional-dependencies]
# vLLM backend (uses OpenAI-compatible API)
vllm = [
"openai>=1.50.0,<2.0",
]
# llama.cpp backend (uses OpenAI-compatible API)
llamacpp = [
"openai>=1.50.0,<2.0",
]
# All backends
all = [
"rerank[vllm,llamacpp]",
]
# Development dependencies
dev = [
"pytest>=8.0",
"pytest-cov>=4.0",
"pytest-asyncio>=0.24.0",
"ruff>=0.8",
"mypy>=1.0",
"respx>=0.21.0",
]
[project.scripts]
rerank = "rerank.cli:main"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
[tool.hatch.build.targets.wheel]
packages = ["src/rerank"]
[tool.ruff]
extend = "../../ruff.toml"
[tool.pytest.ini_options]
asyncio_mode = "auto"
testpaths = ["tests"]
addopts = "-v --tb=short"
[tool.mypy]
python_version = "3.10"
strict = true
warn_return_any = true
warn_unused_ignores = true
[dependency-groups]
dev = [
"pytest>=8.0",
"pytest-asyncio>=0.24.0",
"ruff>=0.8",
]

15
ai_platform/modules/rerank/src/rerank/__init__.py Normal file
View file

@ -0,0 +1,15 @@
"""Rerank - Cohere/Jina-compatible rerank API with multiple backends."""
from rerank.client import RerankClient
from rerank.config import RerankSettings, get_settings
from rerank.types import BackendType
__version__ = "0.1.0"
__all__ = [
"BackendType",
"RerankClient",
"RerankSettings",
"__version__",
"get_settings",
]

1
ai_platform/modules/rerank/src/rerank/api/__init__.py Normal file
View file

@ -0,0 +1 @@
"""FastAPI application for rerank API."""

133
ai_platform/modules/rerank/src/rerank/api/app.py Normal file
View file

@ -0,0 +1,133 @@
"""FastAPI application factory."""
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from fastapi import FastAPI
from rerank.api.dependencies import init_concurrency_limiter
from rerank.api.middleware import RateLimitMiddleware, RequestIdMiddleware
from rerank.api.routes import health, models, rerank
from rerank.client import RerankClient
from rerank.config import SettingsCache
from rerank.logging import configure_logging, get_logger
from rerank.runtime_config import RuntimeConfigClient
logger = get_logger("app")
@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
"""Application lifespan manager."""
settings = SettingsCache.get()
configure_logging(settings.log_level, settings.log_json)
logger.info("Starting Rerank API")
logger.info("Default backend: %s", settings.default_backend)
logger.info("vLLM enabled: %s", settings.enable_vllm)
logger.info("llama.cpp enabled: %s", settings.enable_llamacpp)
init_concurrency_limiter(settings.max_concurrent_reranks)
app.state.settings = settings
app.state.client = RerankClient(settings)
if hasattr(app.state, "runtime_config"):
await app.state.runtime_config.start()
logger.info("Rerank API started successfully")
yield
logger.info("Shutting down Rerank API")
if hasattr(app.state, "runtime_config"):
await app.state.runtime_config.stop()
def create_app() -> FastAPI:
"""Create and configure the FastAPI application."""
settings = SettingsCache.get()
app = FastAPI(
title="Rerank API",
description="Cohere/Jina-compatible reranking with multiple backends (vLLM, llama.cpp)",
version="0.1.0",
lifespan=lifespan,
servers=[{"url": settings.external_url, "description": "Rerank API"}],
)
# ---------------------------------------------------------------- observability
# Prometheus /metrics + OTel tracing (no-op if deps missing or OTEL endpoint unset)
try:
from prometheus_fastapi_instrumentator import Instrumentator as _Inst # type: ignore
_Inst(should_group_status_codes=True).instrument(app).expose(app, endpoint="/metrics", include_in_schema=False)
except ImportError:
pass
import os as _os # noqa: E402
_otel_ep = _os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT")
if _otel_ep:
try:
from opentelemetry import trace as _trace # type: ignore
from opentelemetry.sdk.resources import Resource as _R # type: ignore
from opentelemetry.sdk.trace import TracerProvider as _TP # type: ignore
from opentelemetry.sdk.trace.export import BatchSpanProcessor as _BSP # type: ignore
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter as _Exp # type: ignore
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor as _FInst # type: ignore
_provider = _TP(resource=_R.create({"service.name": _os.environ.get("OTEL_SERVICE_NAME", "didiAI-rerank-api")}))
_provider.add_span_processor(_BSP(_Exp(endpoint=_otel_ep, insecure=True)))
_trace.set_tracer_provider(_provider)
_FInst.instrument_app(app)
print(f"[otel] didiAI-rerank-api instrumented -> {_otel_ep}")
except ImportError as _e:
print(f"[otel] skip: {_e}")
runtime_config = RuntimeConfigClient(
dashboard_url=settings.dashboard_url,
live_log_logger_name="rerank",
live_log_key="rerank.log.level",
)
app.state.runtime_config = runtime_config
app.add_middleware(
RateLimitMiddleware,
rate=settings.rate_limit_rps,
burst=settings.rate_limit_burst,
exclude_paths=["/health", "/ready"],
runtime_config=runtime_config,
rate_key="rerank.rate_limit.rps",
burst_key="rerank.rate_limit.burst",
)
app.add_middleware(RequestIdMiddleware)
app.include_router(health.router, tags=["Health"])
app.include_router(rerank.router, prefix="/v1", tags=["Rerank"])
app.include_router(rerank.router_v2, prefix="/v2", tags=["Rerank"])
app.include_router(models.router, prefix="/v1", tags=["Models"])
return app

159
ai_platform/modules/rerank/src/rerank/api/dependencies.py Normal file
View file

@ -0,0 +1,159 @@
"""FastAPI dependencies for rerank."""
import asyncio
import hmac
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from fastapi import Header, HTTPException, Request
from rerank.client import RerankClient
from rerank.config import RerankSettings
from rerank.logging import get_logger
logger = get_logger("dependencies")
def get_client(request: Request) -> RerankClient:
"""Get the rerank client from application state."""
return request.app.state.client
def get_settings(request: Request) -> RerankSettings:
"""Get settings from application state."""
return request.app.state.settings
def verify_bearer_token(
request: Request,
authorization: str | None = Header(default=None, alias="Authorization"),
) -> str | None:
"""Verify Bearer token authentication.
Returns:
The validated token if auth enabled, None if auth disabled.
Raises:
HTTPException: 401 if auth enabled and token is invalid/missing.
"""
settings = request.app.state.settings
if not settings.auth_enabled:
return None
if not authorization:
raise HTTPException(
status_code=401,
detail={
"error": "Authentication required",
"message": "Missing Authorization header",
},
headers={"WWW-Authenticate": "Bearer"},
)
parts = authorization.split()
if len(parts) != 2 or parts[0].lower() != "bearer":
raise HTTPException(
status_code=401,
detail={
"error": "Authentication required",
"message": "Invalid Authorization header format. Use: Bearer <token>",
},
headers={"WWW-Authenticate": "Bearer"},
)
token = parts[1]
is_valid = any(
hmac.compare_digest(token.encode(), valid_token.encode())
for valid_token in settings.api_tokens
)
if not is_valid:
raise HTTPException(
status_code=401,
detail={
"error": "Authentication failed",
"message": "Invalid API token",
},
headers={"WWW-Authenticate": "Bearer"},
)
return token
class ConcurrencyLimiter:
"""Limits concurrent requests to prevent resource exhaustion."""
def __init__(self, max_concurrent: int) -> None:
"""Initialize concurrency limiter."""
self.max_concurrent = max_concurrent
self._semaphore = asyncio.Semaphore(max_concurrent)
self._counter_lock = asyncio.Lock()
self._current = 0
@property
def current_count(self) -> int:
"""Get current number of active requests."""
return self._current
@property
def available(self) -> int:
"""Get number of available slots."""
return self.max_concurrent - self._current
@asynccontextmanager
async def acquire(self, blocking: bool = False) -> AsyncIterator[None]:
"""Acquire a request slot."""
if not blocking and self._semaphore.locked():
logger.warning(
"Concurrency limit reached: %d/%d active",
self._current,
self.max_concurrent,
)
raise HTTPException(
status_code=503,
detail={
"error": "Service temporarily unavailable",
"reason": "Too many concurrent requests",
"max_concurrent": self.max_concurrent,
"retry_after": 5,
},
headers={"Retry-After": "5"},
)
await self._semaphore.acquire()
async with self._counter_lock:
self._current += 1
try:
yield
finally:
async with self._counter_lock:
self._current -= 1
self._semaphore.release()
_concurrency_limiter: ConcurrencyLimiter | None = None
def init_concurrency_limiter(max_concurrent: int) -> ConcurrencyLimiter:
"""Initialize the global concurrency limiter."""
global _concurrency_limiter
_concurrency_limiter = ConcurrencyLimiter(max_concurrent)
logger.info("Concurrency limiter initialized: max_concurrent=%d", max_concurrent)
return _concurrency_limiter
def get_concurrency_limiter() -> ConcurrencyLimiter:
"""Get the global concurrency limiter."""
if _concurrency_limiter is None:
raise RuntimeError("Concurrency limiter not initialized")
return _concurrency_limiter
async def require_rerank_slot(request: Request) -> AsyncIterator[None]:
"""FastAPI dependency that acquires a rerank slot."""
limiter = get_concurrency_limiter()
async with limiter.acquire():
yield

156
ai_platform/modules/rerank/src/rerank/api/middleware.py Normal file
View file

@ -0,0 +1,156 @@
"""FastAPI middleware for request handling and rate limiting."""
import time
import uuid
from collections.abc import Awaitable, Callable
from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse
from rerank.logging import get_logger, set_request_id
logger = get_logger("middleware")
class RequestIdMiddleware(BaseHTTPMiddleware):
"""Middleware to extract or generate request IDs."""
async def dispatch(
self,
request: Request,
call_next: Callable[[Request], Awaitable[Response]],
) -> Response:
"""Process request with request ID tracking."""
request_id = request.headers.get("X-Request-ID")
if not request_id:
request_id = str(uuid.uuid4())
set_request_id(request_id)
request.state.request_id = request_id
try:
response = await call_next(request)
response.headers["X-Request-ID"] = request_id
return response
finally:
set_request_id(None)
class TokenBucket:
"""Token bucket for rate limiting with optional live tuning."""
def __init__(
self,
rate: float,
burst: int,
runtime_config: object | None = None,
rate_key: str | None = None,
burst_key: str | None = None,
) -> None:
"""Initialize token bucket.
Args:
rate: Tokens added per second (fallback / initial).
burst: Maximum tokens (bucket capacity, fallback / initial).
runtime_config: Optional RuntimeConfigClient for live tuning.
rate_key: Dashboard config key (e.g., "rerank.rate_limit.rps").
burst_key: Dashboard config key (e.g., "rerank.rate_limit.burst").
"""
self._fallback_rate = rate
self._fallback_burst = burst
self.rate = rate
self.burst = burst
self.tokens = float(burst)
self.last_update = time.monotonic()
self._runtime_config = runtime_config
self._rate_key = rate_key
self._burst_key = burst_key
def _refresh_from_config(self) -> None:
"""Pull latest rate/burst from runtime_config if wired."""
if self._runtime_config is None:
return
try:
new_rate = self._runtime_config.get_float(self._rate_key, self._fallback_rate) if self._rate_key else self._fallback_rate
new_burst = self._runtime_config.get_int(self._burst_key, self._fallback_burst) if self._burst_key else self._fallback_burst
except Exception:
return
if new_rate != self.rate or new_burst != self.burst:
self.rate = new_rate
self.burst = new_burst
self.tokens = min(self.tokens, float(new_burst))
def acquire(self) -> bool:
"""Try to acquire a token."""
self._refresh_from_config()
now = time.monotonic()
elapsed = now - self.last_update
self.last_update = now
self.tokens = min(self.burst, self.tokens + elapsed * self.rate)
if self.tokens >= 1:
self.tokens -= 1
return True
return False
def retry_after(self) -> float:
"""Calculate seconds until a token is available."""
if self.tokens >= 1:
return 0.0
tokens_needed = 1 - self.tokens
return tokens_needed / self.rate
class RateLimitMiddleware(BaseHTTPMiddleware):
"""Middleware for global rate limiting using token bucket algorithm."""
def __init__(
self,
app: object,
rate: float = 20.0,
burst: int = 50,
exclude_paths: list[str] | None = None,
runtime_config: object | None = None,
rate_key: str | None = None,
burst_key: str | None = None,
) -> None:
"""Initialize rate limiter."""
super().__init__(app)
self.bucket = TokenBucket(
rate=rate,
burst=burst,
runtime_config=runtime_config,
rate_key=rate_key,
burst_key=burst_key,
)
self.exclude_paths = exclude_paths or ["/health", "/ready"]
async def dispatch(
self,
request: Request,
call_next: Callable[[Request], Awaitable[Response]],
) -> Response:
"""Process request with rate limiting."""
if request.url.path in self.exclude_paths:
return await call_next(request)
if not self.bucket.acquire():
retry_after = self.bucket.retry_after()
logger.warning(
"Rate limit exceeded for %s %s, retry_after=%.2f",
request.method,
request.url.path,
retry_after,
)
return JSONResponse(
status_code=429,
content={
"detail": "Too many requests",
"retry_after": retry_after,
},
headers={"Retry-After": str(int(retry_after) + 1)},
)
return await call_next(request)

1
ai_platform/modules/rerank/src/rerank/api/routes/__init__.py Normal file
View file

@ -0,0 +1 @@
"""API route modules."""

69
ai_platform/modules/rerank/src/rerank/api/routes/health.py Normal file
View file

@ -0,0 +1,69 @@
"""Health check routes."""
from fastapi import APIRouter, Depends
from rerank.api.dependencies import get_client
from rerank.client import RerankClient
from rerank.logging import get_logger
from rerank.schemas import BackendHealth, HealthResponse, ReadinessResponse
router = APIRouter()
logger = get_logger("routes.health")
@router.get("/health", response_model=HealthResponse)
async def health_check(
client: RerankClient = Depends(get_client),
) -> HealthResponse:
"""Health check endpoint.
Returns the overall health status and per-backend health status.
"""
backend_health: list[BackendHealth] = []
for backend_type in client.list_backends():
backend = client.registry.get(backend_type)
is_healthy = await backend.health_check()
backend_health.append(
BackendHealth(
name=backend_type.value,
healthy=is_healthy,
)
)
all_healthy = all(b.healthy for b in backend_health)
any_healthy = any(b.healthy for b in backend_health)
if all_healthy:
status = "healthy"
elif any_healthy:
status = "degraded"
else:
status = "unhealthy"
return HealthResponse(status=status, backends=backend_health)
@router.get("/ready", response_model=ReadinessResponse)
async def readiness_check(
client: RerankClient = Depends(get_client),
) -> ReadinessResponse:
"""Readiness probe for Kubernetes.
Verifies that the default backend is healthy and able to serve requests.
"""
try:
default_backend = client.registry.get()
is_healthy = await default_backend.health_check()
if not is_healthy:
logger.warning(
"Readiness check failed: default backend '%s' is unhealthy",
default_backend.name,
)
return ReadinessResponse(ready=is_healthy)
except Exception as e:
logger.error("Readiness check failed with error: %s", str(e))
return ReadinessResponse(ready=False)

63
ai_platform/modules/rerank/src/rerank/api/routes/models.py Normal file
View file

@ -0,0 +1,63 @@
"""Model management routes."""
from fastapi import APIRouter, Depends, HTTPException
from rerank.api.dependencies import get_client, verify_bearer_token
from rerank.client import RerankClient
from rerank.exceptions import (
BackendNotAvailableError,
BackendNotEnabledError,
ModelListError,
)
from rerank.schemas import BackendListResponse, ModelListResponse
from rerank.types import BackendType
router = APIRouter(dependencies=[Depends(verify_bearer_token)])
@router.get("/models", response_model=ModelListResponse)
async def list_models(
backend: str | None = None,
client: RerankClient = Depends(get_client),
) -> ModelListResponse:
"""List available models.
Args:
backend: Optional backend filter. If not specified, returns models
from all available backends.
client: Rerank client instance.
Returns:
ModelListResponse: List of available models.
Raises:
HTTPException: If the specified backend is not available.
"""
try:
backend_type = BackendType(backend) if backend else None
models = await client.list_models(backend=backend_type)
return ModelListResponse(data=models)
except ValueError:
raise HTTPException(
status_code=400, detail=f"Invalid backend: {backend}"
) from None
except (BackendNotAvailableError, BackendNotEnabledError) as e:
raise HTTPException(status_code=400, detail=str(e)) from None
except ModelListError as e:
raise HTTPException(status_code=503, detail=str(e)) from None
@router.get("/backends", response_model=BackendListResponse)
async def list_backends(
client: RerankClient = Depends(get_client),
) -> BackendListResponse:
"""List available backends.
Args:
client: Rerank client instance.
Returns:
BackendListResponse: List of available backend names.
"""
backends = [b.value for b in client.list_backends()]
return BackendListResponse(backends=backends)

109
ai_platform/modules/rerank/src/rerank/api/routes/rerank.py Normal file
View file

@ -0,0 +1,109 @@
"""Rerank routes."""
from fastapi import APIRouter, Depends, HTTPException
from rerank.api.dependencies import (
get_client,
get_concurrency_limiter,
verify_bearer_token,
)
from rerank.client import RerankClient
from rerank.exceptions import (
BackendNotAvailableError,
BackendNotEnabledError,
RerankConnectionError,
RerankRateLimitError,
RerankRequestError,
RerankTimeoutError,
)
from rerank.logging import get_logger
from rerank.schemas import RerankRequest, RerankResponse
router = APIRouter(dependencies=[Depends(verify_bearer_token)])
router_v2 = APIRouter(dependencies=[Depends(verify_bearer_token)])
logger = get_logger("routes.rerank")
async def _handle_rerank(
request: RerankRequest,
client: RerankClient,
) -> RerankResponse:
"""Common handler for rerank endpoints."""
limiter = get_concurrency_limiter()
async with limiter.acquire():
try:
response = await client.rerank(
query=request.query,
documents=request.documents,
model=request.model,
backend=request.backend,
top_n=request.top_n,
return_documents=request.return_documents,
)
return response
except (BackendNotAvailableError, BackendNotEnabledError) as e:
logger.warning("Backend error: %s", str(e))
raise HTTPException(status_code=400, detail=str(e)) from None
except RerankRateLimitError as e:
logger.warning("Rate limited by backend: %s", str(e))
raise HTTPException(
status_code=429,
detail=str(e),
headers={"Retry-After": str(int(e.retry_after or 5))},
) from None
except RerankTimeoutError as e:
logger.warning("Timeout: %s", str(e))
raise HTTPException(status_code=504, detail=str(e)) from None
except RerankConnectionError as e:
logger.error("Connection error: %s", str(e))
raise HTTPException(status_code=503, detail=str(e)) from None
except RerankRequestError as e:
logger.error("Rerank request failed: %s", str(e))
raise HTTPException(status_code=500, detail=str(e)) from None
except ValueError as e:
logger.warning("Invalid request: %s", str(e))
raise HTTPException(status_code=400, detail=str(e)) from None
@router.post("/rerank", response_model=RerankResponse)
async def rerank_v1(
request: RerankRequest,
client: RerankClient = Depends(get_client),
) -> RerankResponse:
"""Rerank documents against a query.
This endpoint is Cohere/Jina-compatible.
Args:
request: Rerank request with query, documents, and model.
client: Rerank client instance.
Returns:
RerankResponse: The rerank response with scored results.
Raises:
HTTPException: On various error conditions.
"""
return await _handle_rerank(request, client)
@router_v2.post("/rerank", response_model=RerankResponse)
async def rerank_v2(
request: RerankRequest,
client: RerankClient = Depends(get_client),
) -> RerankResponse:
"""Rerank documents against a query (v2 alias).
This endpoint is identical to /v1/rerank for compatibility.
Args:
request: Rerank request with query, documents, and model.
client: Rerank client instance.
Returns:
RerankResponse: The rerank response with scored results.
"""
return await _handle_rerank(request, client)

9
ai_platform/modules/rerank/src/rerank/backends/__init__.py Normal file
View file

@ -0,0 +1,9 @@
"""Backend implementations for reranking."""
from rerank.backends.base import RerankBackend
from rerank.backends.registry import BackendRegistry
__all__ = [
"BackendRegistry",
"RerankBackend",
]

66
ai_platform/modules/rerank/src/rerank/backends/base.py Normal file
View file

@ -0,0 +1,66 @@
"""Abstract base class for rerank backends."""
from abc import ABC, abstractmethod
from rerank.types import ModelInfo, RerankUsage
class RerankBackend(ABC):
"""Abstract base class for rerank backends.
All backend implementations must inherit from this class and implement
the required abstract methods.
"""
@property
@abstractmethod
def name(self) -> str:
"""Backend identifier name.
Returns:
str: Unique name for this backend (e.g., 'vllm', 'llamacpp').
"""
...
@abstractmethod
async def rerank(
self,
query: str,
documents: list[str],
model: str,
top_n: int | None = None,
**kwargs: object,
) -> tuple[list[tuple[int, float]], RerankUsage]:
"""Rerank documents against query.
Args:
query: The search query.
documents: List of documents to rerank.
model: Model identifier to use.
top_n: Number of top results to return.
**kwargs: Additional parameters passed to the backend.
Returns:
Tuple of (list of (original_index, score) sorted by score desc, usage info).
Raises:
RerankRequestError: If the rerank request fails.
"""
...
@abstractmethod
async def list_models(self) -> list[ModelInfo]:
"""List available models for this backend.
Returns:
list[ModelInfo]: List of available models.
"""
...
async def health_check(self) -> bool:
"""Check if this backend is healthy and ready to serve requests.
Returns:
bool: True if the backend is healthy.
"""
return True

220
ai_platform/modules/rerank/src/rerank/backends/llamacpp_backend.py Normal file
View file

@ -0,0 +1,220 @@
"""llama.cpp backend implementation for reranking.
llama.cpp can serve cross-encoder models for reranking via its API.
"""
import logging
import httpx
from rerank.backends.base import RerankBackend
from rerank.config import RerankSettings
from rerank.exceptions import (
ModelListError,
RerankConnectionError,
RerankRateLimitError,
RerankRequestError,
RerankTimeoutError,
)
from rerank.types import ModelInfo, RerankUsage
try:
import openai
from openai import AsyncOpenAI
OPENAI_AVAILABLE = True
except ImportError:
OPENAI_AVAILABLE = False
logger = logging.getLogger("rerank.backends.llamacpp")
class LlamaCppRerankBackend(RerankBackend):
"""llama.cpp backend for reranking.
Supports reranking via llama.cpp's rerank endpoint or by using
the model to score query-document pairs.
Example:
Start llama.cpp server:
```bash
llama-server \
--model /models/bge-reranker-v2-m3-q4_k_m.gguf \
--host 0.0.0.0 --port 54210 \
--reranking
```
Configure:
```bash
export RERANK_ENABLE_LLAMACPP=true
export RERANK_LLAMACPP_BASE_URL=http://localhost:54210
```
"""
def __init__(self, settings: RerankSettings) -> None:
"""Initialize llama.cpp backend.
Args:
settings: Application settings.
Raises:
ImportError: If required packages are not installed.
"""
if not OPENAI_AVAILABLE:
raise ImportError(
"openai package is required for llama.cpp backend. "
"Install with: pip install rerank[llamacpp]"
)
self._settings = settings
timeout = httpx.Timeout(
connect=settings.connect_timeout,
read=settings.request_timeout,
write=settings.request_timeout,
pool=settings.connect_timeout,
)
# For model listing via OpenAI-compatible API
self._openai_client = AsyncOpenAI(
base_url=f"{settings.llamacpp_base_url}/v1",
api_key="not-needed",
timeout=timeout,
max_retries=0,
)
# Direct httpx client for rerank endpoint
self._http_client = httpx.AsyncClient(
base_url=settings.llamacpp_base_url,
timeout=timeout,
)
@property
def name(self) -> str:
"""Backend identifier name."""
return "llamacpp"
async def rerank(
self,
query: str,
documents: list[str],
model: str,
top_n: int | None = None,
**kwargs: object,
) -> tuple[list[tuple[int, float]], RerankUsage]:
"""Rerank documents via llama.cpp.
Args:
query: The search query.
documents: List of documents to rerank.
model: Model identifier (may be ignored by llama.cpp).
top_n: Number of top results to return.
**kwargs: Additional parameters.
Returns:
Tuple of (list of (index, score), usage).
Raises:
RerankRequestError: If the request fails.
"""
try:
# Try llama.cpp's /rerank endpoint
response = await self._http_client.post(
"/rerank",
json={
"query": query,
"documents": documents,
"top_k": top_n or len(documents),
},
)
response.raise_for_status()
data = response.json()
# Extract results - llama.cpp returns results array
results = []
for r in data.get("results", []):
idx = r.get("index", 0)
score = r.get("relevance_score", r.get("score", 0.0))
results.append((idx, score))
# Sort by score descending
results.sort(key=lambda x: x[1], reverse=True)
# Apply top_n if specified
if top_n is not None:
results = results[:top_n]
usage = RerankUsage(
total_tokens=data.get("usage", {}).get("total_tokens", 0),
)
return results, usage
except httpx.HTTPStatusError as e:
if e.response.status_code == 429:
logger.warning("Rate limited: %s", str(e))
raise RerankRateLimitError(str(e), backend=self.name) from e
logger.error("HTTP error: %s", str(e))
raise RerankRequestError(str(e), backend=self.name, model=model) from e
except httpx.TimeoutException as e:
logger.warning("Timeout: %s", str(e))
raise RerankTimeoutError(str(e), backend=self.name) from e
except httpx.ConnectError as e:
logger.error("Connection error: %s", str(e))
raise RerankConnectionError(self.name, reason=str(e)) from e
except Exception as e:
logger.error("Unexpected error: %s", str(e))
raise RerankRequestError(str(e), backend=self.name, model=model) from e
async def list_models(self) -> list[ModelInfo]:
"""List models available on the llama.cpp server.
Note: llama.cpp typically serves a single model.
Returns:
list[ModelInfo]: List of available models.
Raises:
ModelListError: If models cannot be retrieved.
"""
try:
models = await self._openai_client.models.list()
return [
ModelInfo(
id=m.id,
backend=self.name,
)
for m in models.data
]
except openai.APIConnectionError as e:
logger.error("Failed to list models - connection error: %s", str(e))
raise ModelListError(
self.name, f"Cannot connect to llama.cpp server: {e}"
) from e
except openai.APITimeoutError as e:
logger.error("Failed to list models - timeout: %s", str(e))
raise ModelListError(
self.name, f"Timeout connecting to llama.cpp server: {e}"
) from e
except Exception as e:
logger.error("Failed to list models: %s", str(e))
raise ModelListError(self.name, str(e)) from e
async def health_check(self) -> bool:
"""Check if llama.cpp server is healthy.
Returns:
bool: True if the server is reachable and responding.
"""
try:
await self._openai_client.models.list()
return True
except openai.APIConnectionError:
logger.debug("llama.cpp health check failed: connection error")
return False
except openai.APITimeoutError:
logger.debug("llama.cpp health check failed: timeout")
return False
except Exception as e:
logger.debug("llama.cpp health check failed: %s", str(e))
return False

106
ai_platform/modules/rerank/src/rerank/backends/registry.py Normal file
View file

@ -0,0 +1,106 @@
"""Backend registry for managing rerank backend instances."""
from rerank.backends.base import RerankBackend
from rerank.config import RerankSettings
from rerank.exceptions import BackendNotAvailableError, BackendNotEnabledError
from rerank.types import BackendType
class BackendRegistry:
"""Registry for managing rerank backend instances.
The registry initializes and provides access to configured backends.
Both vLLM and llama.cpp are optional and require explicit enabling.
"""
def __init__(self, settings: RerankSettings) -> None:
"""Initialize the backend registry.
Args:
settings: Application settings.
"""
self._settings = settings
self._backends: dict[BackendType, RerankBackend] = {}
self._initialize_backends()
def _initialize_backends(self) -> None:
"""Initialize enabled backends."""
# vLLM is optional but MUST work if enabled (fail-fast)
if self._settings.enable_vllm:
try:
from rerank.backends.vllm_backend import VLLMRerankBackend
self._backends[BackendType.VLLM] = VLLMRerankBackend(self._settings)
except ImportError as e:
raise BackendNotAvailableError(
"vllm",
f"vLLM backend is enabled but dependencies are not installed. "
f"Install with: pip install rerank[vllm]. "
f"Original error: {e}",
) from e
# llama.cpp is optional but MUST work if enabled (fail-fast)
if self._settings.enable_llamacpp:
try:
from rerank.backends.llamacpp_backend import LlamaCppRerankBackend
self._backends[BackendType.LLAMACPP] = LlamaCppRerankBackend(
self._settings
)
except ImportError as e:
raise BackendNotAvailableError(
"llamacpp",
f"llama.cpp backend is enabled but dependencies are not installed. "
f"Install with: pip install rerank[llamacpp]. "
f"Original error: {e}",
) from e
# Validate that at least one backend is enabled
if not self._backends:
raise BackendNotAvailableError(
"none",
"No backends are enabled. Enable at least one backend: "
"RERANK_ENABLE_VLLM=true or RERANK_ENABLE_LLAMACPP=true",
)
def get(self, backend_type: BackendType | str | None = None) -> RerankBackend:
"""Get a backend instance.
Args:
backend_type: Backend to retrieve. If None, uses default from settings.
Returns:
RerankBackend: The requested backend instance.
Raises:
BackendNotAvailableError: If the backend is not available.
BackendNotEnabledError: If the backend is not enabled.
"""
if backend_type is None:
backend_type = BackendType(self._settings.default_backend)
elif isinstance(backend_type, str):
backend_type = BackendType(backend_type)
if backend_type not in self._backends:
if backend_type == BackendType.VLLM and not self._settings.enable_vllm:
raise BackendNotEnabledError("vllm")
if (
backend_type == BackendType.LLAMACPP
and not self._settings.enable_llamacpp
):
raise BackendNotEnabledError("llamacpp")
raise BackendNotAvailableError(
backend_type.value, "Backend not initialized"
)
return self._backends[backend_type]
def list_backends(self) -> list[BackendType]:
"""List all available backends."""
return list(self._backends.keys())
def is_available(self, backend_type: BackendType | str) -> bool:
"""Check if a backend is available."""
if isinstance(backend_type, str):
backend_type = BackendType(backend_type)
return backend_type in self._backends

210
ai_platform/modules/rerank/src/rerank/backends/vllm_backend.py Normal file
View file

@ -0,0 +1,210 @@
"""vLLM backend implementation for reranking.
vLLM supports reranking via its /score or /rerank endpoint when started
with --task score for cross-encoder models.
"""
import logging
import httpx
from rerank.backends.base import RerankBackend
from rerank.config import RerankSettings
from rerank.exceptions import (
ModelListError,
RerankConnectionError,
RerankRateLimitError,
RerankRequestError,
RerankTimeoutError,
)
from rerank.types import ModelInfo, RerankUsage
try:
import openai
from openai import AsyncOpenAI
OPENAI_AVAILABLE = True
except ImportError:
OPENAI_AVAILABLE = False
logger = logging.getLogger("rerank.backends.vllm")
class VLLMRerankBackend(RerankBackend):
"""vLLM backend for reranking using cross-encoder models.
Requires the vLLM server to be running with --task score.
Example:
Start vLLM server:
```bash
python -m vllm.entrypoints.openai.api_server \
--model BAAI/bge-reranker-v2-m3 \
--host 0.0.0.0 --port 54201 \
--task score
```
Configure:
```bash
export RERANK_ENABLE_VLLM=true
export RERANK_VLLM_BASE_URL=http://localhost:54201
```
"""
def __init__(self, settings: RerankSettings) -> None:
"""Initialize vLLM backend.
Args:
settings: Application settings.
Raises:
ImportError: If openai/httpx packages are not installed.
"""
if not OPENAI_AVAILABLE:
raise ImportError(
"openai package is required for vLLM backend. "
"Install with: pip install rerank[vllm]"
)
self._settings = settings
timeout = httpx.Timeout(
connect=settings.connect_timeout,
read=settings.request_timeout,
write=settings.request_timeout,
pool=settings.connect_timeout,
)
# For model listing via OpenAI-compatible API
self._openai_client = AsyncOpenAI(
base_url=f"{settings.vllm_base_url}/v1",
api_key=settings.vllm_api_key or "not-needed",
timeout=timeout,
max_retries=0,
)
# Direct httpx client for rerank endpoint
self._http_client = httpx.AsyncClient(
base_url=settings.vllm_base_url,
timeout=timeout,
)
@property
def name(self) -> str:
"""Backend identifier name."""
return "vllm"
async def rerank(
self,
query: str,
documents: list[str],
model: str,
top_n: int | None = None,
**kwargs: object,
) -> tuple[list[tuple[int, float]], RerankUsage]:
"""Rerank documents via vLLM's /rerank endpoint.
Args:
query: The search query.
documents: List of documents to rerank.
model: Model identifier.
top_n: Number of top results to return.
**kwargs: Additional parameters.
Returns:
Tuple of (list of (index, score), usage).
Raises:
RerankRequestError: If the request fails.
"""
try:
# vLLM's rerank endpoint (Jina-compatible)
response = await self._http_client.post(
"/rerank",
json={
"model": model,
"query": query,
"documents": documents,
"top_n": top_n or len(documents),
},
)
response.raise_for_status()
data = response.json()
# Extract results
results = [(r["index"], r["relevance_score"]) for r in data["results"]]
# Extract usage if available
usage_data = data.get("usage", {})
usage = RerankUsage(
total_tokens=usage_data.get("total_tokens", 0),
)
return results, usage
except httpx.HTTPStatusError as e:
if e.response.status_code == 429:
logger.warning("Rate limited: %s", str(e))
raise RerankRateLimitError(str(e), backend=self.name) from e
logger.error("HTTP error: %s", str(e))
raise RerankRequestError(str(e), backend=self.name, model=model) from e
except httpx.TimeoutException as e:
logger.warning("Timeout: %s", str(e))
raise RerankTimeoutError(str(e), backend=self.name) from e
except httpx.ConnectError as e:
logger.error("Connection error: %s", str(e))
raise RerankConnectionError(self.name, reason=str(e)) from e
except Exception as e:
logger.error("Unexpected error: %s", str(e))
raise RerankRequestError(str(e), backend=self.name, model=model) from e
async def list_models(self) -> list[ModelInfo]:
"""List models available on the vLLM server.
Returns:
list[ModelInfo]: List of available models.
Raises:
ModelListError: If models cannot be retrieved.
"""
try:
models = await self._openai_client.models.list()
return [
ModelInfo(
id=m.id,
backend=self.name,
)
for m in models.data
]
except openai.APIConnectionError as e:
logger.error("Failed to list models - connection error: %s", str(e))
raise ModelListError(
self.name, f"Cannot connect to vLLM server: {e}"
) from e
except openai.APITimeoutError as e:
logger.error("Failed to list models - timeout: %s", str(e))
raise ModelListError(
self.name, f"Timeout connecting to vLLM server: {e}"
) from e
except Exception as e:
logger.error("Failed to list models: %s", str(e))
raise ModelListError(self.name, str(e)) from e
async def health_check(self) -> bool:
"""Check if vLLM server is healthy.
Returns:
bool: True if the server is reachable and responding.
"""
try:
await self._openai_client.models.list()
return True
except openai.APIConnectionError:
logger.debug("vLLM health check failed: connection error")
return False
except openai.APITimeoutError:
logger.debug("vLLM health check failed: timeout")
return False
except Exception as e:
logger.debug("vLLM health check failed: %s", str(e))
return False

123
ai_platform/modules/rerank/src/rerank/cli.py Normal file
View file

@ -0,0 +1,123 @@
"""CLI entry point for the rerank server."""
import argparse
import signal
import sys
from typing import Any
import uvicorn
from rerank.config import SettingsCache
from rerank.logging import configure_logging, get_logger
class GracefulShutdown:
"""Handles graceful shutdown on signals."""
def __init__(self) -> None:
"""Initialize shutdown handler."""
self._server: uvicorn.Server | None = None
self._shutdown_requested = False
self._logger = get_logger("cli.shutdown")
def register_signals(self) -> None:
"""Register signal handlers for graceful shutdown."""
signal.signal(signal.SIGTERM, self._signal_handler)
signal.signal(signal.SIGINT, self._signal_handler)
def _signal_handler(self, signum: int, frame: Any) -> None:
"""Handle shutdown signals."""
if self._shutdown_requested:
self._logger.warning("Forced shutdown requested")
sys.exit(1)
self._shutdown_requested = True
signal_name = signal.Signals(signum).name
self._logger.info("Received %s, initiating graceful shutdown...", signal_name)
if self._server:
self._server.should_exit = True
def set_server(self, server: uvicorn.Server) -> None:
"""Set the uvicorn server for shutdown control."""
self._server = server
def main() -> None:
"""Run the rerank server with graceful shutdown support."""
parser = argparse.ArgumentParser(
description="Rerank API Server",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument(
"--host",
type=str,
default=None,
help="Host to bind to (default: from RERANK_HOST env or 0.0.0.0)",
)
parser.add_argument(
"--port",
type=int,
default=None,
help="Port to bind to (default: from RERANK_PORT env or 54200)",
)
parser.add_argument(
"--workers",
type=int,
default=1,
help="Number of worker processes",
)
parser.add_argument(
"--reload",
action="store_true",
help="Enable auto-reload for development",
)
parser.add_argument(
"--graceful-timeout",
type=int,
default=30,
help="Timeout in seconds for graceful shutdown",
)
args = parser.parse_args()
settings = SettingsCache.get()
configure_logging(settings.log_level, settings.log_json)
logger = get_logger("cli")
host = args.host or settings.host
port = args.port or settings.port
logger.info("Starting Rerank API Server on %s:%d", host, port)
if args.reload:
logger.info("Development mode - auto-reload enabled")
uvicorn.run(
"rerank.api.app:create_app",
factory=True,
host=host,
port=port,
workers=1,
reload=True,
)
else:
shutdown_handler = GracefulShutdown()
shutdown_handler.register_signals()
config = uvicorn.Config(
"rerank.api.app:create_app",
factory=True,
host=host,
port=port,
workers=args.workers,
timeout_graceful_shutdown=args.graceful_timeout,
)
server = uvicorn.Server(config)
shutdown_handler.set_server(server)
server.run()
if __name__ == "__main__":
main()

161
ai_platform/modules/rerank/src/rerank/client.py Normal file
View file

@ -0,0 +1,161 @@
"""High-level rerank client for unified inference."""
from rerank.backends import BackendRegistry
from rerank.config import RerankSettings, SettingsCache
from rerank.schemas import RerankResponse
from rerank.types import BackendType, ModelInfo, RerankResult
class RerankClient:
"""High-level client for reranking.
This is the primary interface for using the module as a Python library.
It provides a unified interface across all configured backends.
Example:
```python
from rerank import RerankClient
client = RerankClient()
# Rerank documents
response = await client.rerank(
query="What is machine learning?",
documents=[
"ML is a subset of AI",
"Cats are pets",
"Deep learning uses neural networks",
],
model="BAAI/bge-reranker-v2-m3",
)
for result in response.results:
print(f"Index {result.index}: {result.relevance_score:.4f}")
# With top_n
response = await client.rerank(
query="programming",
documents=["Python", "Java", "Cooking", "C++"],
model="BAAI/bge-reranker-v2-m3",
top_n=2,
)
# Override backend per-request
response = await client.rerank(
query="test",
documents=["doc1", "doc2"],
model="bge-reranker-v2-m3.gguf",
backend=BackendType.LLAMACPP,
)
```
"""
def __init__(self, settings: RerankSettings | None = None) -> None:
"""Initialize the rerank client.
Args:
settings: Optional settings instance. If not provided, loads from environment.
"""
self._settings = settings or SettingsCache.get()
self._registry = BackendRegistry(self._settings)
@property
def registry(self) -> BackendRegistry:
"""Get the backend registry."""
return self._registry
async def rerank(
self,
query: str,
documents: list[str],
model: str,
backend: BackendType | str | None = None,
top_n: int | None = None,
return_documents: bool = False,
**kwargs: object,
) -> RerankResponse:
"""Rerank documents against a query.
Args:
query: The search query.
documents: List of documents to rerank.
model: Model identifier.
backend: Backend to use. Uses default if not specified.
top_n: Number of top results to return (default: all).
return_documents: Include document text in response.
**kwargs: Additional parameters.
Returns:
RerankResponse: The rerank response.
Raises:
BackendNotAvailableError: If the requested backend is not available.
RerankRequestError: If the rerank request fails.
"""
backend_instance = self._registry.get(backend)
results_data, usage = await backend_instance.rerank(
query=query,
documents=documents,
model=model,
top_n=top_n,
**kwargs,
)
# Build response
results = [
RerankResult(
index=idx,
relevance_score=score,
document=documents[idx] if return_documents else None,
)
for idx, score in results_data
]
return RerankResponse(
model=model,
results=results,
usage=usage,
backend=backend_instance.name,
)
async def list_models(
self,
backend: BackendType | str | None = None,
) -> list[ModelInfo]:
"""List available models.
Args:
backend: If specified, only list models for this backend.
Otherwise, lists models from all available backends.
Returns:
list[ModelInfo]: List of available models.
"""
if backend is not None:
return await self._registry.get(backend).list_models()
# Aggregate from all backends
all_models: list[ModelInfo] = []
for backend_type in self._registry.list_backends():
backend_instance = self._registry.get(backend_type)
models = await backend_instance.list_models()
all_models.extend(models)
return all_models
def list_backends(self) -> list[BackendType]:
"""List available backends."""
return self._registry.list_backends()
async def health_check(self) -> dict[str, bool]:
"""Check health of all backends.
Returns:
dict[str, bool]: Mapping of backend name to health status.
"""
health: dict[str, bool] = {}
for backend_type in self._registry.list_backends():
backend_instance = self._registry.get(backend_type)
health[backend_type.value] = await backend_instance.health_check()
return health

200
ai_platform/modules/rerank/src/rerank/config.py Normal file
View file

@ -0,0 +1,200 @@
"""Configuration management using Pydantic V2 Settings."""
import threading
from typing import Literal
from pydantic import Field, field_validator
from pydantic_settings import BaseSettings, SettingsConfigDict
class RerankSettings(BaseSettings):
"""Rerank module configuration.
All settings can be configured via environment variables with the RERANK_ prefix.
Required environment variables (no defaults - fail-fast):
RERANK_DEFAULT_BACKEND: Backend to use (vllm, llamacpp)
RERANK_ENABLE_VLLM: Enable vLLM backend (true/false)
RERANK_ENABLE_LLAMACPP: Enable llama.cpp backend (true/false)
RERANK_EXTERNAL_URL: External URL for OpenAPI spec
Example:
RERANK_DEFAULT_BACKEND=vllm
RERANK_ENABLE_VLLM=true
RERANK_ENABLE_LLAMACPP=false
RERANK_EXTERNAL_URL=http://localhost:54200
"""
model_config = SettingsConfigDict(
env_prefix="RERANK_",
env_file=".env",
env_file_encoding="utf-8",
extra="forbid",
)
# ==========================================================================
# REQUIRED fields (no defaults - fail-fast)
# ==========================================================================
default_backend: Literal["vllm", "llamacpp"] = Field(
description="Default backend to use for reranking (REQUIRED)",
)
enable_vllm: bool = Field(
description="Enable vLLM backend (REQUIRED, set to true or false)",
)
enable_llamacpp: bool = Field(
description="Enable llama.cpp backend (REQUIRED, set to true or false)",
)
external_url: str = Field(
description="External URL for OpenAPI spec (REQUIRED)",
)
# ==========================================================================
# Optional fields with sensible defaults
# ==========================================================================
# API server settings
host: str = Field(
default="0.0.0.0",
description="Host to bind the server to",
)
port: int = Field(
default=54200,
description="Port to bind the server to (Dev: 54200, Prod: 14200)",
)
# vLLM settings
vllm_base_url: str = Field(
default="http://localhost:54201",
description="Base URL for vLLM rerank server",
)
vllm_api_key: str | None = Field(
default=None,
description="API key for vLLM server (if required)",
)
# llama.cpp settings
llamacpp_base_url: str = Field(
default="http://localhost:54210",
description="Base URL for llama.cpp rerank server",
)
# ==========================================================================
# Timeout settings
# ==========================================================================
request_timeout: float = Field(
default=120.0,
ge=1.0,
description="Default timeout in seconds for rerank requests",
)
connect_timeout: float = Field(
default=10.0,
ge=1.0,
description="Timeout in seconds for establishing connections",
)
# ==========================================================================
# Observability settings
# ==========================================================================
log_level: str = Field(
default="INFO",
description="Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)",
)
log_json: bool = Field(
default=False,
description="Output logs in JSON format for production",
)
# ==========================================================================
# Rate limiting and concurrency
# ==========================================================================
rate_limit_rps: float = Field(
default=20.0,
ge=0.1,
description="Rate limit: requests per second",
)
rate_limit_burst: int = Field(
default=50,
ge=1,
description="Rate limit: burst capacity",
)
max_concurrent_reranks: int = Field(
default=20,
ge=1,
description="Maximum concurrent rerank requests",
)
# ==========================================================================
# Authentication settings
# ==========================================================================
api_tokens: frozenset[str] | None = Field(
default=None,
description=(
"API tokens for Bearer authentication (comma-separated). "
"If not set, authentication is disabled."
),
)
# ==========================================================================
# Runtime config (dashboard polling)
# ==========================================================================
dashboard_url: str | None = Field(
default=None,
description=(
"Optional dashboard base URL. When set, runtime_config polls "
"/api/config every 30s for live overrides."
),
)
@field_validator("api_tokens", mode="before")
@classmethod
def parse_api_tokens(cls, v: str | list[str] | None) -> frozenset[str] | None:
"""Parse comma-separated tokens into a frozenset."""
if v is None or v == "":
return None
if isinstance(v, str):
tokens = [t.strip() for t in v.split(",") if t.strip()]
return frozenset(tokens) if tokens else None
return frozenset(v) if v else None
@property
def auth_enabled(self) -> bool:
"""Check if authentication is enabled."""
return bool(self.api_tokens)
class SettingsCache:
"""Thread-safe settings cache that can be cleared for testing."""
_instance: RerankSettings | None = None
_lock: threading.Lock = threading.Lock()
@classmethod
def get(cls) -> RerankSettings:
"""Get or create the settings instance."""
with cls._lock:
if cls._instance is None:
cls._instance = RerankSettings()
return cls._instance
@classmethod
def clear(cls) -> None:
"""Clear the cached settings instance."""
with cls._lock:
cls._instance = None
@classmethod
def set(cls, settings: RerankSettings) -> None:
"""Set a specific settings instance."""
with cls._lock:
cls._instance = settings
def get_settings() -> RerankSettings:
"""Get cached settings instance."""
return SettingsCache.get()

124
ai_platform/modules/rerank/src/rerank/exceptions.py Normal file
View file

@ -0,0 +1,124 @@
"""Custom exceptions for rerank module."""
class RerankError(Exception):
"""Base exception for rerank errors."""
class AuthenticationError(RerankError):
"""Raised when authentication fails."""
def __init__(self, message: str = "Authentication required") -> None:
super().__init__(message)
class BackendNotAvailableError(RerankError):
"""Raised when a requested backend is not available."""
def __init__(self, backend: str, reason: str | None = None) -> None:
self.backend = backend
self.reason = reason
message = f"Backend '{backend}' is not available"
if reason:
message += f": {reason}"
super().__init__(message)
class BackendNotEnabledError(RerankError):
"""Raised when a backend is not enabled in configuration."""
def __init__(self, backend: str) -> None:
self.backend = backend
super().__init__(
f"Backend '{backend}' is not enabled. "
f"Set RERANK_ENABLE_{backend.upper()}=true to enable it."
)
class ModelNotFoundError(RerankError):
"""Raised when a requested model is not found."""
def __init__(self, model: str, backend: str | None = None) -> None:
self.model = model
self.backend = backend
message = f"Model '{model}' not found"
if backend:
message += f" on backend '{backend}'"
super().__init__(message)
class RerankRequestError(RerankError):
"""Raised when a rerank request fails."""
def __init__(
self,
message: str,
backend: str | None = None,
model: str | None = None,
) -> None:
self.backend = backend
self.model = model
full_message = message
if backend or model:
details = []
if backend:
details.append(f"backend={backend}")
if model:
details.append(f"model={model}")
full_message = f"{message} ({', '.join(details)})"
super().__init__(full_message)
class ModelListError(RerankError):
"""Raised when listing models fails."""
def __init__(self, backend: str, reason: str | None = None) -> None:
self.backend = backend
self.reason = reason
message = f"Failed to list models on backend '{backend}'"
if reason:
message += f": {reason}"
super().__init__(message)
class RerankRateLimitError(RerankError):
"""Raised when a rate limit is hit."""
def __init__(
self,
message: str,
backend: str | None = None,
retry_after: float | None = None,
) -> None:
self.backend = backend
self.retry_after = retry_after
full_message = message
if retry_after:
full_message += f" (retry after {retry_after}s)"
super().__init__(full_message)
class RerankTimeoutError(RerankError):
"""Raised when a request times out."""
def __init__(
self,
message: str,
backend: str | None = None,
timeout: float | None = None,
) -> None:
self.backend = backend
self.timeout = timeout
super().__init__(message)
class RerankConnectionError(RerankError):
"""Raised when connection to backend fails."""
def __init__(self, backend: str, reason: str | None = None) -> None:
self.backend = backend
self.reason = reason
message = f"Connection to backend '{backend}' failed"
if reason:
message += f": {reason}"
super().__init__(message)

85
ai_platform/modules/rerank/src/rerank/logging.py Normal file
View file

@ -0,0 +1,85 @@
"""Structured logging with request_id context propagation."""
import logging
import sys
from contextvars import ContextVar
from typing import Any
request_id_ctx: ContextVar[str | None] = ContextVar("request_id", default=None)
class RequestIdFilter(logging.Filter):
"""Logging filter that adds request_id to log records."""
def filter(self, record: logging.LogRecord) -> bool:
"""Add request_id to the log record."""
record.request_id = request_id_ctx.get() or "-"
return True
class JsonFormatter(logging.Formatter):
"""JSON log formatter for structured logging."""
def format(self, record: logging.LogRecord) -> str:
"""Format log record as JSON."""
import json
log_data: dict[str, Any] = {
"timestamp": self.formatTime(record, self.datefmt),
"level": record.levelname,
"logger": record.name,
"message": record.getMessage(),
"request_id": getattr(record, "request_id", "-"),
}
if record.exc_info:
log_data["exception"] = self.formatException(record.exc_info)
if hasattr(record, "extra_data"):
log_data.update(record.extra_data)
return json.dumps(log_data)
def configure_logging(level: str = "INFO", json_format: bool = False) -> None:
"""Configure logging for the application."""
root_logger = logging.getLogger("rerank")
root_logger.setLevel(getattr(logging, level.upper(), logging.INFO))
for handler in root_logger.handlers[:]:
root_logger.removeHandler(handler)
handler = logging.StreamHandler(sys.stdout)
handler.setLevel(logging.DEBUG)
handler.addFilter(RequestIdFilter())
if json_format:
handler.setFormatter(JsonFormatter())
else:
handler.setFormatter(
logging.Formatter(
"%(asctime)s [%(levelname)s] [%(request_id)s] %(name)s: %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
)
root_logger.addHandler(handler)
root_logger.propagate = False
def get_logger(name: str) -> logging.Logger:
"""Get a logger with the rerank prefix."""
if name.startswith("rerank."):
return logging.getLogger(name)
return logging.getLogger(f"rerank.{name}")
def set_request_id(request_id: str | None) -> None:
"""Set the request ID for the current context."""
request_id_ctx.set(request_id)
def get_request_id() -> str | None:
"""Get the request ID for the current context."""
return request_id_ctx.get()

157
ai_platform/modules/rerank/src/rerank/runtime_config.py Normal file
View file

@ -0,0 +1,157 @@
"""Runtime config client — fetches config overrides from dashboard.
Polls the dashboard /api/config endpoint periodically and caches values
in memory. All accessors fall back to caller-supplied defaults if the
dashboard is unreachable or the key is missing.
Single source of truth = dashboard `KNOWN_KEYS`. This client carries no
local default registry; consumers pass their own fallback (typically the
pydantic settings value).
"""
import asyncio
import contextlib
import logging
from typing import Any
import httpx
logger = logging.getLogger(__name__)
class RuntimeConfigClient:
"""Polls dashboard /api/config and caches values in-process."""
def __init__(
self,
dashboard_url: str | None,
poll_interval_seconds: int = 30,
request_timeout: float = 5.0,
live_log_logger_name: str | None = None,
live_log_key: str | None = None,
) -> None:
self.dashboard_url = dashboard_url.rstrip("/") if dashboard_url else None
self.poll_interval = poll_interval_seconds
self._timeout = request_timeout
self._cache: dict[str, Any] = {}
self._client: httpx.AsyncClient | None = None
self._task: asyncio.Task[None] | None = None
self._enabled = bool(dashboard_url)
# Optional auto-apply: re-set log level on the named logger when the
# configured key changes value (e.g., `llm.log.level`).
self._live_log_logger_name = live_log_logger_name
self._live_log_key = live_log_key
self._last_log_level: str | None = None
@property
def enabled(self) -> bool:
return self._enabled
def get(self, key: str, default: Any = None) -> Any:
v = self._cache.get(key)
return v if v is not None else default
def get_bool(self, key: str, default: bool = False) -> bool:
v = self._cache.get(key)
return bool(v) if v is not None else default
def get_int(self, key: str, default: int = 0) -> int:
v = self._cache.get(key)
if v is None:
return default
try:
return int(v)
except (TypeError, ValueError):
return default
def get_float(self, key: str, default: float = 0.0) -> float:
v = self._cache.get(key)
if v is None:
return default
try:
return float(v)
except (TypeError, ValueError):
return default
def get_str(self, key: str, default: str = "") -> str:
v = self._cache.get(key)
return str(v) if v is not None else default
async def start(self) -> None:
if not self._enabled:
logger.info("RuntimeConfigClient disabled (no dashboard URL)")
return
self._client = httpx.AsyncClient(
timeout=httpx.Timeout(
connect=2.0, read=self._timeout, write=2.0, pool=5.0
)
)
# Fetch once synchronously so first requests already have overrides
await self._refresh()
self._task = asyncio.create_task(self._loop())
logger.info(
"RuntimeConfigClient started (polling %s every %ds, %d keys cached)",
self.dashboard_url,
self.poll_interval,
len(self._cache),
)
async def stop(self) -> None:
if self._task is not None:
self._task.cancel()
with contextlib.suppress(asyncio.CancelledError, Exception):
await self._task
self._task = None
if self._client is not None and not self._client.is_closed:
await self._client.aclose()
self._client = None
async def _loop(self) -> None:
while True:
try:
await asyncio.sleep(self.poll_interval)
await self._refresh()
except asyncio.CancelledError:
raise
except Exception as e:
logger.debug("Config poll failed: %s", e)
async def _refresh(self) -> None:
if self._client is None or self.dashboard_url is None:
return
try:
resp = await self._client.get(f"{self.dashboard_url}/api/config")
resp.raise_for_status()
data = resp.json()
except Exception as e:
logger.debug("Config refresh failed: %s", e)
return
items = data.get("items", {})
new_cache: dict[str, Any] = {}
for key, entry in items.items():
new_cache[key] = entry.get("value")
self._cache = new_cache
self._maybe_apply_log_level()
def _maybe_apply_log_level(self) -> None:
"""Re-apply log level live if configured key changed."""
if not self._live_log_key or not self._live_log_logger_name:
return
new_level = self.get_str(self._live_log_key)
if not new_level:
return
if new_level == self._last_log_level:
return
try:
level_int = logging.getLevelName(new_level.upper())
if isinstance(level_int, int):
logging.getLogger(self._live_log_logger_name).setLevel(level_int)
self._last_log_level = new_level
logger.info(
"Log level for %s changed to %s (via runtime config)",
self._live_log_logger_name,
new_level,
)
except Exception as e:
logger.warning("Failed to apply log level %s: %s", new_level, e)

115
ai_platform/modules/rerank/src/rerank/schemas.py Normal file
View file

@ -0,0 +1,115 @@
"""Pydantic V2 API request/response schemas."""
import uuid
from pydantic import BaseModel, ConfigDict, Field, field_validator
from rerank.types import BackendType, ModelInfo, RerankResult, RerankUsage
# =============================================================================
# Rerank Schemas
# =============================================================================
class RerankRequest(BaseModel):
"""Request schema for reranking (Cohere/Jina-compatible)."""
model_config = ConfigDict(extra="forbid")
model: str = Field(
description="Model to use for reranking",
)
query: str = Field(
min_length=1,
description="The search query",
)
documents: list[str] = Field(
min_length=1,
max_length=1000,
description="List of documents to rerank (1-1000)",
)
top_n: int | None = Field(
default=None,
ge=1,
description="Number of top results to return (default: all)",
)
return_documents: bool = Field(
default=False,
description="Include document text in response",
)
backend: BackendType | None = Field(
default=None,
description="Backend to use (overrides default)",
)
@field_validator("documents")
@classmethod
def validate_documents(cls, v: list[str]) -> list[str]:
"""Validate documents are not empty strings."""
if any(not doc.strip() for doc in v):
raise ValueError("documents cannot contain empty strings")
return v
class RerankResponse(BaseModel):
"""Response schema for reranking (Cohere/Jina-compatible)."""
model_config = ConfigDict(extra="forbid")
id: str = Field(
default_factory=lambda: f"rerank-{uuid.uuid4().hex[:12]}",
description="Unique request ID",
)
model: str = Field(description="Model used for reranking")
results: list[RerankResult] = Field(description="Reranked results")
usage: RerankUsage = Field(description="Token usage information")
backend: str = Field(description="Backend that served the request")
# =============================================================================
# Model Management Schemas
# =============================================================================
class ModelListResponse(BaseModel):
"""Response schema for listing models."""
object: str = Field(default="list")
data: list[ModelInfo] = Field(description="List of available models")
# =============================================================================
# Health Schemas
# =============================================================================
class BackendHealth(BaseModel):
"""Health status for a single backend."""
name: str = Field(description="Backend name")
healthy: bool = Field(description="Whether the backend is healthy")
message: str | None = Field(default=None, description="Optional status message")
class HealthResponse(BaseModel):
"""Response schema for health check."""
status: str = Field(description="Overall status: healthy, degraded, unhealthy")
backends: list[BackendHealth] = Field(description="Per-backend health status")
class ReadinessResponse(BaseModel):
"""Response schema for readiness probe."""
ready: bool = Field(default=True)
# =============================================================================
# Backend Schemas
# =============================================================================
class BackendListResponse(BaseModel):
"""Response schema for listing backends."""
backends: list[str] = Field(description="List of available backend names")

43
ai_platform/modules/rerank/src/rerank/types.py Normal file
View file

@ -0,0 +1,43 @@
"""Core types and enums for reranking."""
from enum import Enum
from pydantic import BaseModel, ConfigDict, Field
class BackendType(str, Enum):
"""Supported reranking backends."""
VLLM = "vllm"
LLAMACPP = "llamacpp"
class RerankUsage(BaseModel):
"""Token usage information for reranking."""
total_tokens: int = Field(default=0)
class RerankResult(BaseModel):
"""A single rerank result."""
model_config = ConfigDict(extra="forbid")
index: int = Field(description="Original index of the document")
relevance_score: float = Field(description="Relevance score")
document: str | None = Field(
default=None,
description="Document text (if return_documents=True)",
)
class ModelInfo(BaseModel):
"""Information about an available model."""
id: str = Field(description="Model identifier")
backend: str = Field(description="Backend serving this model")
loaded: bool = Field(default=True, description="Whether the model is loaded")
max_input_tokens: int | None = Field(
default=None,
description="Maximum input tokens",
)

1
ai_platform/modules/rerank/tests/__init__.py Normal file
View file

@ -0,0 +1 @@
"""Tests for rerank module."""

81
ai_platform/modules/rerank/tests/conftest.py Normal file
View file

@ -0,0 +1,81 @@
"""Pytest fixtures for rerank tests."""
from unittest.mock import AsyncMock
import pytest
from fastapi.testclient import TestClient
from rerank.api.app import create_app
from rerank.api.dependencies import init_concurrency_limiter
from rerank.client import RerankClient
from rerank.config import RerankSettings, SettingsCache
from rerank.schemas import RerankResponse
from rerank.types import BackendType, RerankResult, RerankUsage
@pytest.fixture(autouse=True)
def clear_settings_cache() -> None:
"""Clear settings cache before each test."""
SettingsCache.clear()
@pytest.fixture
def test_settings() -> RerankSettings:
"""Create test settings with mocked values."""
return RerankSettings(
default_backend="vllm",
enable_vllm=True,
enable_llamacpp=False,
external_url="http://localhost:54200",
host="127.0.0.1",
port=54200,
api_tokens=None,
)
@pytest.fixture
def mock_rerank_response() -> RerankResponse:
"""Create a mock rerank response."""
return RerankResponse(
model="test-model",
results=[
RerankResult(index=2, relevance_score=0.95),
RerankResult(index=0, relevance_score=0.82),
RerankResult(index=1, relevance_score=0.10),
],
usage=RerankUsage(total_tokens=150),
backend="vllm",
)
@pytest.fixture
def client_with_mock_backend(
test_settings: RerankSettings, mock_rerank_response: RerankResponse
) -> RerankClient:
"""Create a RerankClient with mocked backend."""
SettingsCache.set(test_settings)
client = RerankClient(settings=test_settings)
# Mock the backend's rerank method
backend = client.registry.get(BackendType.VLLM)
results = [(r.index, r.relevance_score) for r in mock_rerank_response.results]
usage = mock_rerank_response.usage
backend.rerank = AsyncMock(return_value=(results, usage))
return client
@pytest.fixture
def app_client(test_settings: RerankSettings) -> TestClient:
"""Create a FastAPI TestClient with test settings."""
SettingsCache.set(test_settings)
init_concurrency_limiter(test_settings.max_concurrent_reranks)
app = create_app()
app.state.settings = test_settings
app.state.client = RerankClient(settings=test_settings)
return TestClient(app)

179
ai_platform/modules/rerank/tests/test_config.py Normal file
View file

@ -0,0 +1,179 @@
"""Tests for configuration module."""
import os
from unittest.mock import patch
import pytest
from pydantic import ValidationError
from rerank.config import RerankSettings, SettingsCache
class TestRerankSettings:
"""Tests for RerankSettings."""
def test_required_fields(self) -> None:
"""Test that required fields must be provided."""
with pytest.raises(ValidationError):
RerankSettings()
def test_with_required_fields(self) -> None:
"""Test settings with all required fields provided."""
settings = RerankSettings(
default_backend="vllm",
enable_vllm=True,
enable_llamacpp=False,
external_url="http://localhost:54200",
)
assert settings.default_backend == "vllm"
assert settings.enable_vllm is True
assert settings.enable_llamacpp is False
assert settings.external_url == "http://localhost:54200"
def test_optional_defaults(self) -> None:
"""Test optional fields have sensible defaults."""
settings = RerankSettings(
default_backend="vllm",
enable_vllm=True,
enable_llamacpp=False,
external_url="http://localhost:54200",
)
assert settings.host == "0.0.0.0"
assert settings.port == 54200
assert settings.request_timeout == 120.0
assert settings.rate_limit_rps == 20.0
assert settings.max_concurrent_reranks == 20
def test_env_override(self) -> None:
"""Test environment variable overrides."""
env_vars = {
"RERANK_DEFAULT_BACKEND": "llamacpp",
"RERANK_ENABLE_VLLM": "false",
"RERANK_ENABLE_LLAMACPP": "true",
"RERANK_EXTERNAL_URL": "http://example.com",
"RERANK_PORT": "14200",
}
with patch.dict(os.environ, env_vars, clear=False):
settings = RerankSettings()
assert settings.default_backend == "llamacpp"
assert settings.enable_vllm is False
assert settings.enable_llamacpp is True
assert settings.port == 14200
def test_vllm_url_default(self) -> None:
"""Test vLLM URL default value."""
settings = RerankSettings(
default_backend="vllm",
enable_vllm=True,
enable_llamacpp=False,
external_url="http://localhost:54200",
)
assert settings.vllm_base_url == "http://localhost:54201"
def test_llamacpp_url_default(self) -> None:
"""Test llama.cpp URL default value."""
settings = RerankSettings(
default_backend="vllm",
enable_vllm=True,
enable_llamacpp=False,
external_url="http://localhost:54200",
)
assert settings.llamacpp_base_url == "http://localhost:54210"
def test_api_tokens_parsing(self) -> None:
"""Test API tokens are parsed from comma-separated string."""
settings = RerankSettings(
default_backend="vllm",
enable_vllm=True,
enable_llamacpp=False,
external_url="http://localhost:54200",
api_tokens="token1,token2,token3", # type: ignore[arg-type]
)
assert settings.api_tokens is not None
assert "token1" in settings.api_tokens
assert "token2" in settings.api_tokens
assert "token3" in settings.api_tokens
assert settings.auth_enabled is True
def test_auth_disabled_by_default(self) -> None:
"""Test authentication is disabled when no tokens set."""
settings = RerankSettings(
default_backend="vllm",
enable_vllm=True,
enable_llamacpp=False,
external_url="http://localhost:54200",
)
assert settings.api_tokens is None
assert settings.auth_enabled is False
class TestSettingsCache:
"""Tests for SettingsCache class."""
def test_set_and_get(self) -> None:
"""Test setting and getting cached settings."""
SettingsCache.clear()
settings = RerankSettings(
default_backend="vllm",
enable_vllm=True,
enable_llamacpp=False,
external_url="http://localhost:54200",
)
SettingsCache.set(settings)
retrieved = SettingsCache.get()
assert retrieved is settings
def test_get_without_set_creates_instance(self) -> None:
"""Test that get() creates instance from environment if not set."""
SettingsCache.clear()
env_vars = {
"RERANK_DEFAULT_BACKEND": "vllm",
"RERANK_ENABLE_VLLM": "true",
"RERANK_ENABLE_LLAMACPP": "false",
"RERANK_EXTERNAL_URL": "http://localhost:54200",
}
with patch.dict(os.environ, env_vars, clear=False):
settings = SettingsCache.get()
assert isinstance(settings, RerankSettings)
assert settings.default_backend == "vllm"
def test_cached_returns_same_instance(self) -> None:
"""Test that SettingsCache returns cached instance."""
SettingsCache.clear()
settings = RerankSettings(
default_backend="vllm",
enable_vllm=True,
enable_llamacpp=False,
external_url="http://localhost:54200",
)
SettingsCache.set(settings)
settings1 = SettingsCache.get()
settings2 = SettingsCache.get()
assert settings1 is settings2
def test_clear_removes_cached_instance(self) -> None:
"""Test that clear() removes the cached instance."""
settings = RerankSettings(
default_backend="vllm",
enable_vllm=True,
enable_llamacpp=False,
external_url="http://localhost:54200",
)
SettingsCache.set(settings)
SettingsCache.clear()
assert SettingsCache._instance is None

145
ai_platform/modules/rerank/tests/test_schemas.py Normal file
View file

@ -0,0 +1,145 @@
"""Tests for schemas module."""
import pytest
from pydantic import ValidationError
from rerank.schemas import RerankRequest, RerankResponse
from rerank.types import BackendType, RerankResult, RerankUsage
class TestRerankRequest:
"""Tests for RerankRequest schema."""
def test_basic_request(self) -> None:
"""Test basic request creation."""
request = RerankRequest(
model="test-model",
query="What is AI?",
documents=["doc1", "doc2", "doc3"],
)
assert request.model == "test-model"
assert request.query == "What is AI?"
assert len(request.documents) == 3
def test_empty_documents_rejected(self) -> None:
"""Test that empty documents list is rejected."""
with pytest.raises(ValidationError):
RerankRequest(
model="test-model",
query="query",
documents=[],
)
def test_empty_string_in_documents_rejected(self) -> None:
"""Test that empty strings in documents are rejected."""
with pytest.raises(ValidationError):
RerankRequest(
model="test-model",
query="query",
documents=["doc1", "", "doc3"],
)
def test_empty_query_rejected(self) -> None:
"""Test that empty query is rejected."""
with pytest.raises(ValidationError):
RerankRequest(
model="test-model",
query="",
documents=["doc1", "doc2"],
)
def test_top_n_parameter(self) -> None:
"""Test top_n parameter."""
request = RerankRequest(
model="test-model",
query="query",
documents=["doc1", "doc2", "doc3"],
top_n=2,
)
assert request.top_n == 2
def test_return_documents_parameter(self) -> None:
"""Test return_documents parameter."""
request = RerankRequest(
model="test-model",
query="query",
documents=["doc1", "doc2"],
return_documents=True,
)
assert request.return_documents is True
def test_backend_override(self) -> None:
"""Test backend override."""
request = RerankRequest(
model="test-model",
query="query",
documents=["doc1", "doc2"],
backend=BackendType.LLAMACPP,
)
assert request.backend == BackendType.LLAMACPP
def test_invalid_top_n(self) -> None:
"""Test that invalid top_n is rejected."""
with pytest.raises(ValidationError):
RerankRequest(
model="test-model",
query="query",
documents=["doc1", "doc2"],
top_n=0,
)
class TestRerankResponse:
"""Tests for RerankResponse schema."""
def test_creation(self) -> None:
"""Test creating response."""
response = RerankResponse(
model="test-model",
results=[
RerankResult(index=2, relevance_score=0.95),
RerankResult(index=0, relevance_score=0.82),
],
usage=RerankUsage(total_tokens=100),
backend="vllm",
)
assert response.model == "test-model"
assert len(response.results) == 2
assert response.backend == "vllm"
assert response.id.startswith("rerank-")
def test_results_ordering(self) -> None:
"""Test that results maintain order."""
results = [
RerankResult(index=2, relevance_score=0.95),
RerankResult(index=0, relevance_score=0.82),
RerankResult(index=1, relevance_score=0.10),
]
response = RerankResponse(
model="test-model",
results=results,
usage=RerankUsage(total_tokens=100),
backend="vllm",
)
assert response.results[0].index == 2
assert response.results[0].relevance_score == 0.95
assert response.results[1].index == 0
assert response.results[2].index == 1
class TestRerankResult:
"""Tests for RerankResult model."""
def test_without_document(self) -> None:
"""Test result without document."""
result = RerankResult(index=0, relevance_score=0.95)
assert result.document is None
def test_with_document(self) -> None:
"""Test result with document."""
result = RerankResult(
index=0,
relevance_score=0.95,
document="This is the document text",
)
assert result.document == "This is the document text"