Livrare LOT 1 - Didi
This commit is contained in:
commit
5380c3fc63
990 changed files with 133308 additions and 0 deletions
123
ai_platform/modules/embeddings/.env.example
Normal file
123
ai_platform/modules/embeddings/.env.example
Normal file
|
|
@ -0,0 +1,123 @@
|
|||
# Embeddings 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
|
||||
EMB_DEFAULT_BACKEND=vllm
|
||||
|
||||
# Enable backends (true/false)
|
||||
EMB_ENABLE_VLLM=true
|
||||
EMB_ENABLE_LLAMACPP=false
|
||||
|
||||
# External URL for OpenAPI spec (REQUIRED)
|
||||
EMB_EXTERNAL_URL=http://localhost:14100
|
||||
|
||||
# =============================================================================
|
||||
# OPTIONAL: API Server Settings
|
||||
# =============================================================================
|
||||
|
||||
# Server binding
|
||||
# EMB_HOST=0.0.0.0
|
||||
# EMB_PORT=14100 # Prod: 14100, Dev: 54100
|
||||
|
||||
# =============================================================================
|
||||
# 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
|
||||
# EMB_API_TOKENS=token1,token2,token3
|
||||
|
||||
# =============================================================================
|
||||
# OPTIONAL: vLLM Backend Configuration
|
||||
# =============================================================================
|
||||
|
||||
# vLLM server URL (internal Docker network)
|
||||
# EMB_VLLM_BASE_URL=http://localhost:54101
|
||||
|
||||
# vLLM model to load (Hugging Face model ID)
|
||||
EMB_VLLM_MODEL=BAAI/bge-m3
|
||||
|
||||
# vLLM server port
|
||||
# EMB_VLLM_PORT=14101
|
||||
|
||||
# GPU assignment for vLLM
|
||||
# EMB_VLLM_GPU=0
|
||||
|
||||
# GPU memory utilization
|
||||
# EMB_VLLM_GPU_UTIL=0.50
|
||||
|
||||
# Maximum model sequence length
|
||||
# EMB_VLLM_MAX_LEN=8192
|
||||
|
||||
# =============================================================================
|
||||
# OPTIONAL: llama.cpp Backend Configuration
|
||||
# =============================================================================
|
||||
|
||||
# llama.cpp server URL (internal Docker network)
|
||||
# EMB_LLAMACPP_BASE_URL=http://localhost:54110
|
||||
|
||||
# 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)
|
||||
EMB_LLAMACPP_MODEL=bge-m3-q4_k_m.gguf
|
||||
|
||||
# llama.cpp server port
|
||||
# EMB_LLAMACPP_PORT=14110
|
||||
|
||||
# llama.cpp context size
|
||||
# EMB_LLAMACPP_CTX=8192
|
||||
|
||||
# llama.cpp performance settings
|
||||
# EMB_LLAMACPP_THREADS=4
|
||||
# EMB_LLAMACPP_PARALLEL=4
|
||||
|
||||
# =============================================================================
|
||||
# OPTIONAL: Timeout Settings
|
||||
# =============================================================================
|
||||
|
||||
# Request timeout for embedding calls (seconds)
|
||||
# EMB_REQUEST_TIMEOUT=120.0
|
||||
|
||||
# Connection timeout (seconds)
|
||||
# EMB_CONNECT_TIMEOUT=10.0
|
||||
|
||||
# =============================================================================
|
||||
# OPTIONAL: Rate Limiting and Concurrency
|
||||
# =============================================================================
|
||||
|
||||
# Requests per second limit
|
||||
# EMB_RATE_LIMIT_RPS=20.0
|
||||
|
||||
# Maximum burst size for rate limiting
|
||||
# EMB_RATE_LIMIT_BURST=40
|
||||
|
||||
# Maximum concurrent embedding requests
|
||||
# EMB_MAX_CONCURRENT_REQUESTS=20
|
||||
|
||||
# =============================================================================
|
||||
# OPTIONAL: Logging
|
||||
# =============================================================================
|
||||
|
||||
# Log level: DEBUG, INFO, WARNING, ERROR, CRITICAL
|
||||
# EMB_LOG_LEVEL=INFO
|
||||
|
||||
# Enable JSON logging format (true/false)
|
||||
# EMB_LOG_JSON=false
|
||||
|
||||
# =============================================================================
|
||||
# OPTIONAL: Shared Settings (for Docker)
|
||||
# =============================================================================
|
||||
|
||||
# HuggingFace cache directory
|
||||
# HF_CACHE_DIR=/cai2_ds_storage/hf_cache
|
||||
|
||||
# HuggingFace token
|
||||
# HF_TOKEN=
|
||||
336
ai_platform/modules/embeddings/API.md
Normal file
336
ai_platform/modules/embeddings/API.md
Normal file
|
|
@ -0,0 +1,336 @@
|
|||
# Embeddings API Reference
|
||||
|
||||
OpenAI-compatible embeddings API with multiple backend support.
|
||||
|
||||
## Base URL
|
||||
|
||||
```
|
||||
{BASE_URL}
|
||||
```
|
||||
|
||||
- **Local development:** `http://localhost:54100`
|
||||
- **Docker (internal):** `http://didiAI-embeddings-api:14100`
|
||||
- **Production:** Use your configured hostname
|
||||
|
||||
## Authentication
|
||||
|
||||
Authentication is **optional**. If `EMB_API_TOKENS` is set, requests require a Bearer token:
|
||||
|
||||
```
|
||||
Authorization: Bearer <token>
|
||||
```
|
||||
|
||||
Health endpoints (`/health`, `/ready`) are always public.
|
||||
|
||||
## Endpoints
|
||||
|
||||
### Create Embeddings
|
||||
|
||||
Generate embeddings for the given input texts.
|
||||
|
||||
**Endpoint:** `POST /v1/embeddings`
|
||||
|
||||
**Request Headers:**
|
||||
| Header | Required | Description |
|
||||
|--------|----------|-------------|
|
||||
| `Content-Type` | Yes | Must be `application/json` |
|
||||
| `Authorization` | If auth enabled | `Bearer <token>` |
|
||||
|
||||
**Request Body:**
|
||||
```json
|
||||
{
|
||||
"input": "text to embed",
|
||||
"model": "BAAI/bge-m3",
|
||||
"encoding_format": "float",
|
||||
"dimensions": null,
|
||||
"backend": null
|
||||
}
|
||||
```
|
||||
|
||||
| Field | Type | Required | Description |
|
||||
|-------|------|----------|-------------|
|
||||
| `input` | string or string[] | Yes | Text(s) to embed |
|
||||
| `model` | string | Yes | Model identifier |
|
||||
| `encoding_format` | string | No | `"float"` (default) or `"base64"` |
|
||||
| `dimensions` | integer | No | Desired embedding dimensions (if supported) |
|
||||
| `backend` | string | No | Override default backend: `"vllm"` or `"llamacpp"` |
|
||||
|
||||
**Response:**
|
||||
```json
|
||||
{
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"object": "embedding",
|
||||
"index": 0,
|
||||
"embedding": [0.0023, -0.0142, 0.0083, ...]
|
||||
}
|
||||
],
|
||||
"model": "BAAI/bge-m3",
|
||||
"usage": {
|
||||
"prompt_tokens": 5,
|
||||
"total_tokens": 5
|
||||
},
|
||||
"backend": "vllm"
|
||||
}
|
||||
```
|
||||
|
||||
**Example:**
|
||||
```bash
|
||||
curl -X POST http://localhost:54100/v1/embeddings \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": ["Hello world", "How are you?"],
|
||||
"model": "BAAI/bge-m3"
|
||||
}'
|
||||
```
|
||||
|
||||
**Multiple texts:**
|
||||
```bash
|
||||
curl -X POST http://localhost:54100/v1/embeddings \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": [
|
||||
"First document to embed",
|
||||
"Second document to embed",
|
||||
"Third document to embed"
|
||||
],
|
||||
"model": "BAAI/bge-m3"
|
||||
}'
|
||||
```
|
||||
|
||||
**Base64 encoding:**
|
||||
```bash
|
||||
curl -X POST http://localhost:54100/v1/embeddings \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello world",
|
||||
"model": "BAAI/bge-m3",
|
||||
"encoding_format": "base64"
|
||||
}'
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### List Models
|
||||
|
||||
List available embedding 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-m3",
|
||||
"backend": "vllm",
|
||||
"loaded": true,
|
||||
"dimensions": null,
|
||||
"max_input_tokens": null
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
**Example:**
|
||||
```bash
|
||||
# List all models
|
||||
curl http://localhost:54100/v1/models
|
||||
|
||||
# List models from specific backend
|
||||
curl "http://localhost:54100/v1/models?backend=vllm"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### List Backends
|
||||
|
||||
List available backends.
|
||||
|
||||
**Endpoint:** `GET /v1/backends`
|
||||
|
||||
**Response:**
|
||||
```json
|
||||
{
|
||||
"backends": ["vllm", "llamacpp"]
|
||||
}
|
||||
```
|
||||
|
||||
**Example:**
|
||||
```bash
|
||||
curl http://localhost:54100/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:54100/health
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Readiness Probe
|
||||
|
||||
Simple readiness check for Kubernetes.
|
||||
|
||||
**Endpoint:** `GET /ready`
|
||||
|
||||
**Response:**
|
||||
```json
|
||||
{
|
||||
"ready": true
|
||||
}
|
||||
```
|
||||
|
||||
**Example:**
|
||||
```bash
|
||||
curl http://localhost:54100/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`
|
||||
|
||||
### Authentication Error
|
||||
|
||||
```json
|
||||
{
|
||||
"detail": {
|
||||
"error": "Authentication required",
|
||||
"message": "Missing Authorization header"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 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 embed_texts(texts: list[str]) -> list[list[float]]:
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.post(
|
||||
"http://localhost:54100/v1/embeddings",
|
||||
json={
|
||||
"input": texts,
|
||||
"model": "BAAI/bge-m3",
|
||||
},
|
||||
headers={"Authorization": "Bearer your-token"},
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
return [item["embedding"] for item in data["data"]]
|
||||
```
|
||||
|
||||
### Python (openai SDK)
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(
|
||||
base_url="http://localhost:54100/v1",
|
||||
api_key="your-token", # or "not-needed" if auth disabled
|
||||
)
|
||||
|
||||
response = client.embeddings.create(
|
||||
input=["Hello world"],
|
||||
model="BAAI/bge-m3",
|
||||
)
|
||||
|
||||
embedding = response.data[0].embedding
|
||||
print(f"Embedding dimensions: {len(embedding)}")
|
||||
```
|
||||
|
||||
### curl
|
||||
|
||||
```bash
|
||||
# Simple embedding
|
||||
curl -X POST http://localhost:54100/v1/embeddings \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer your-token" \
|
||||
-d '{"input": "Hello world", "model": "BAAI/bge-m3"}'
|
||||
|
||||
# With specific backend
|
||||
curl -X POST http://localhost:54100/v1/embeddings \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"input": "Hello world", "model": "BAAI/bge-m3", "backend": "vllm"}'
|
||||
```
|
||||
214
ai_platform/modules/embeddings/INDEX.md
Normal file
214
ai_platform/modules/embeddings/INDEX.md
Normal file
|
|
@ -0,0 +1,214 @@
|
|||
# Embeddings Module — INDEX
|
||||
|
||||
Embedding generation API for DIDI knowledge graph + semantic search. Serves `BAAI/bge-m3` (1024-dim, multilingual, 8K context) via vLLM (or llama.cpp) behind an OpenAI-compatible FastAPI wrapper, used by `didi-brain` for atom storage in pgvector and for `/v1/gather` query embedding.
|
||||
|
||||
- **Stack**: Python 3.10+, FastAPI, uvicorn, httpx, Pydantic v2; backends via `vllm/vllm-openai:v0.8.5` or `ghcr.io/ggml-org/llama.cpp:server`
|
||||
- **URLs**:
|
||||
- API wrapper: `http://10.11.10.12:14100` (Prod) / `http://10.11.10.12:54100` (Dev)
|
||||
- vLLM internal: `:14101` (Prod) / `:54101` (Dev)
|
||||
- llama.cpp internal: `:14110` (Prod) / `:54110` (Dev)
|
||||
- **Container layout (2-tier)**:
|
||||
- `didiAI-embeddings-api` (FastAPI wrapper + auth + rate limiting)
|
||||
- `didiAI-embeddings-vllm` (GPU embedding server) **OR** `didiAI-embeddings-llamacpp` (CPU/lightweight server)
|
||||
- **Default model**: `BAAI/bge-m3` — 1024 dim, multilingual, 8K context
|
||||
|
||||
---
|
||||
|
||||
## Ce face
|
||||
|
||||
Generate text embeddings for downstream DIDI consumers:
|
||||
|
||||
- **didi-brain semantic atom storage** — text claims/atoms get embedded once at ingestion, stored as pgvector columns.
|
||||
- **didi-brain `/v1/gather` query embedding** — incoming claim text gets embedded, used as the kNN probe vector.
|
||||
- **Cross-encoder reranking** — separate `rerank` module (port 54200) handles re-scoring; this module only does dense embeddings.
|
||||
- **Catalog / retrieval consumers** — anything that needs a vector representation hits this single endpoint.
|
||||
|
||||
The module is OpenAI-compatible: drop-in for `openai` SDK code that uses `client.embeddings.create()`.
|
||||
|
||||
---
|
||||
|
||||
## API endpoints
|
||||
|
||||
All routes are mounted by `src/embeddings/api/app.py` via `create_app()`. Health is unauthenticated and excluded from rate-limiting; everything else passes through `verify_bearer_token` if `EMB_API_TOKENS` is set.
|
||||
|
||||
| Method | Path | Description |
|
||||
|--------|------|-------------|
|
||||
| `POST` | `/v1/embeddings` | OpenAI-compatible. Body: `{input, model, encoding_format?, dimensions?, backend?}`. Returns `{object, data[], model, usage, backend}`. |
|
||||
| `GET` | `/v1/models` | List available embedding models. Optional `?backend=vllm\|llamacpp` filter. |
|
||||
| `GET` | `/v1/backends` | List enabled backends as `{backends: ["vllm", ...]}`. |
|
||||
| `GET` | `/health` | Detailed health: per-backend status, overall `healthy/degraded/unhealthy`. |
|
||||
| `GET` | `/ready` | Simple readiness probe (`{ready: true}`) — for k8s. |
|
||||
|
||||
Source files:
|
||||
- `/home/admin365/didi_mono/ai_platform/modules/embeddings/src/embeddings/api/routes/embeddings.py` — `POST /v1/embeddings`
|
||||
- `/home/admin365/didi_mono/ai_platform/modules/embeddings/src/embeddings/api/routes/models.py` — `/v1/models`, `/v1/backends`
|
||||
- `/home/admin365/didi_mono/ai_platform/modules/embeddings/src/embeddings/api/routes/health.py` — `/health`, `/ready`
|
||||
|
||||
Errors: 400 (invalid params / backend disabled), 401 (auth), 429 (rate limit, sets `Retry-After`), 503 (backend connection), 504 (backend timeout), 500 (other backend failures).
|
||||
|
||||
---
|
||||
|
||||
## Backends
|
||||
|
||||
Two interchangeable engines, selected per-request via the `backend` field or globally via `EMB_DEFAULT_BACKEND`:
|
||||
|
||||
- **vLLM (production / GPU)** — `vllm/vllm-openai:v0.8.5` launched with `--task embed`, high-throughput batching, NVIDIA CUDA 12.x required. Listens internally on `14101` (Prod) / `54101` (Dev).
|
||||
- **llama.cpp (dev / lightweight)** — `ghcr.io/ggml-org/llama.cpp:server-b4769` launched with `--embedding`, runs on CPU or GPU with GGUF-quantized models. Listens on `14110` / `54110`.
|
||||
- **API wrapper** — own FastAPI app on `14100` (Prod) / `54100` (Dev) that proxies to whichever backend(s) are enabled.
|
||||
|
||||
Both backends speak OpenAI-compatible HTTP, so the wrapper just routes via `httpx`/`openai` SDK. Backend abstraction lives in `src/embeddings/backends/{base,vllm_backend,llamacpp_backend,registry}.py`.
|
||||
|
||||
---
|
||||
|
||||
## Structura fisiere
|
||||
|
||||
```
|
||||
embeddings/
|
||||
├── INDEX.md (this file)
|
||||
├── README.md quick-start + config table
|
||||
├── API.md full HTTP API reference + SDK examples
|
||||
├── pyproject.toml package def, optional extras: vllm, llamacpp, all, dev
|
||||
├── uv.lock
|
||||
├── .env.example all EMB_* env vars documented
|
||||
├── deploy/
|
||||
│ ├── Dockerfile multi-stage, python 3.11-slim + uv
|
||||
│ ├── docker-compose.yml 3 services + 3 profiles (api/vllm/llamacpp)
|
||||
│ └── deploy.sh wrapper around `docker compose --profile`
|
||||
├── src/embeddings/
|
||||
│ ├── __init__.py re-exports EmbeddingClient
|
||||
│ ├── cli.py `python -m embeddings.cli` entry point
|
||||
│ ├── client.py EmbeddingClient (Python library API)
|
||||
│ ├── config.py EmbeddingSettings (pydantic-settings, EMB_ prefix)
|
||||
│ ├── schemas.py EmbeddingRequest/Response/Data, base64 encoder
|
||||
│ ├── types.py BackendType enum, ModelInfo
|
||||
│ ├── exceptions.py BackendNotAvailable/Enabled, RateLimit, Timeout, etc.
|
||||
│ ├── logging.py structured logging (text or JSON)
|
||||
│ ├── api/
|
||||
│ │ ├── app.py create_app() factory + lifespan
|
||||
│ │ ├── dependencies.py FastAPI deps (auth, concurrency limiter, client)
|
||||
│ │ ├── middleware.py RateLimitMiddleware, RequestIdMiddleware
|
||||
│ │ └── routes/
|
||||
│ │ ├── embeddings.py POST /v1/embeddings
|
||||
│ │ ├── models.py /v1/models, /v1/backends
|
||||
│ │ └── health.py /health, /ready
|
||||
│ └── backends/
|
||||
│ ├── base.py abstract Backend interface
|
||||
│ ├── vllm_backend.py vLLM HTTP client
|
||||
│ ├── llamacpp_backend.py llama.cpp HTTP client
|
||||
│ └── registry.py BackendRegistry (selects + caches backends)
|
||||
└── tests/
|
||||
├── conftest.py
|
||||
├── test_config.py
|
||||
├── test_schemas.py
|
||||
└── test_types.py
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## How didi-brain uses it
|
||||
|
||||
`didi-brain` consumes this service through `shared/embedding_client.py` (a thin wrapper around the `openai` SDK pointed at this API):
|
||||
|
||||
- **Ingestion path** — when atoms/claims are written, brain calls `EmbeddingClient.embed(text)` to get a 1024-dim vector and stores it in pgvector alongside the row. One round-trip per batch.
|
||||
- **Query path** — for `/v1/gather`, the incoming claim text is embedded the same way, then the resulting vector is used as the probe in a pgvector `<=> ` (cosine) kNN search to retrieve candidate atoms.
|
||||
- **Reranking** — top-k candidates from the dense search are forwarded to the separate `rerank` module (cross-encoder, port 54200) for fine-grained scoring. That module does not call this one — they're parallel concerns.
|
||||
- **Catalog API** — also retrieves via the same embedding pipeline (text → vector → kNN), reusing this single endpoint.
|
||||
|
||||
Because the wrapper is OpenAI-compatible, `EmbeddingClient` can be a vanilla `openai.OpenAI(base_url=..., api_key=...)` instance — no DIDI-specific client code needed in brain.
|
||||
|
||||
---
|
||||
|
||||
## Configuration
|
||||
|
||||
All env vars use the `EMB_` prefix. Required vars have no defaults — the app refuses to start without them.
|
||||
|
||||
### Required
|
||||
|
||||
| Variable | Description |
|
||||
|----------|-------------|
|
||||
| `EMB_DEFAULT_BACKEND` | `vllm` or `llamacpp` |
|
||||
| `EMB_ENABLE_VLLM` | `true` / `false` |
|
||||
| `EMB_ENABLE_LLAMACPP` | `true` / `false` |
|
||||
| `EMB_EXTERNAL_URL` | Public URL exposed in the OpenAPI spec |
|
||||
|
||||
### Common optional
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `EMB_PORT` | `14100` (Prod) / `54100` (Dev) | API server port |
|
||||
| `EMB_HOST` | `0.0.0.0` | API bind address |
|
||||
| `EMB_API_TOKENS` | unset | Comma-separated bearer tokens; auth disabled if unset |
|
||||
| `EMB_VLLM_BASE_URL` | `http://localhost:54101` | Where vLLM listens |
|
||||
| `EMB_VLLM_MODEL` | `BAAI/bge-m3` | HF model id loaded by vLLM |
|
||||
| `EMB_VLLM_GPU` | `0` | `CUDA_VISIBLE_DEVICES` for vLLM |
|
||||
| `EMB_VLLM_GPU_UTIL` | `0.50` | vLLM `--gpu-memory-utilization` |
|
||||
| `EMB_VLLM_MAX_LEN` | `8192` | vLLM `--max-model-len` |
|
||||
| `EMB_LLAMACPP_BASE_URL` | `http://localhost:54110` | Where llama.cpp listens |
|
||||
| `EMB_LLAMACPP_MODEL` | `bge-m3-q4_k_m.gguf` | GGUF filename inside `MODELS_DIR` |
|
||||
| `EMB_LLAMACPP_CTX` | `8192` | llama.cpp context size |
|
||||
| `EMB_LLAMACPP_THREADS` | `4` | llama.cpp threads |
|
||||
| `EMB_LLAMACPP_PARALLEL` | `4` | llama.cpp parallel slots |
|
||||
| `EMB_REQUEST_TIMEOUT` | `120.0` | Per-request backend timeout (s) |
|
||||
| `EMB_CONNECT_TIMEOUT` | `10.0` | TCP connect timeout (s) |
|
||||
| `EMB_RATE_LIMIT_RPS` | `20.0` | Token-bucket rate (req/s) |
|
||||
| `EMB_RATE_LIMIT_BURST` | `40` | Token-bucket burst capacity |
|
||||
| `EMB_MAX_CONCURRENT_REQUESTS` | `20` | In-flight cap |
|
||||
| `EMB_LOG_LEVEL` | `INFO` | DEBUG / INFO / WARNING / ERROR |
|
||||
| `EMB_LOG_JSON` | `false` | Emit JSON-formatted log lines |
|
||||
| `HF_CACHE_DIR` | `/cai2_ds_storage/hf_cache` | Mounted into vLLM container |
|
||||
| `HF_TOKEN` | unset | Forwarded as `HUGGING_FACE_HUB_TOKEN` |
|
||||
| `MODELS_DIR` | `/cai2_ds_storage/models` | GGUF model directory for llama.cpp |
|
||||
|
||||
See `/home/admin365/didi_mono/ai_platform/modules/embeddings/.env.example` for the canonical, fully commented list.
|
||||
|
||||
---
|
||||
|
||||
## Deployment
|
||||
|
||||
GPU host required for the production profile (e.g. `10.11.10.17`). All three services live on the shared `didi-network` Docker network.
|
||||
|
||||
```bash
|
||||
cd /home/admin365/didi_mono/ai_platform/modules/embeddings/deploy
|
||||
|
||||
# Configure
|
||||
cp ../.env.example .env
|
||||
$EDITOR .env
|
||||
|
||||
# Production: API + vLLM (GPU)
|
||||
./deploy.sh --profile vllm -d
|
||||
|
||||
# Lightweight: API + llama.cpp
|
||||
./deploy.sh --profile llamacpp -d
|
||||
|
||||
# API only (use external embedding servers via EMB_*_BASE_URL)
|
||||
./deploy.sh --profile api -d
|
||||
|
||||
# Logs / shutdown
|
||||
./deploy.sh --profile vllm --logs
|
||||
./deploy.sh --profile vllm --down
|
||||
```
|
||||
|
||||
Compose profiles:
|
||||
- `api` — only the FastAPI wrapper
|
||||
- `vllm` — wrapper + vLLM (GPU)
|
||||
- `llamacpp` — wrapper + llama.cpp (CPU/GPU)
|
||||
|
||||
---
|
||||
|
||||
## Performance / model
|
||||
|
||||
- **Model**: `BAAI/bge-m3` — 1024 dim, multilingual (100+ languages), 8K token context, supports dense + sparse + multi-vector (this module uses dense only).
|
||||
- **Throughput**: ~100–200 embeddings/s on H100 with vLLM batching (depends on input length).
|
||||
- **Latency P95**: ~50–100 ms per request (single text, warm GPU). Larger batches amortize well — keep request batches at 16–64 inputs for best throughput.
|
||||
- **Concurrency knobs**: `EMB_MAX_CONCURRENT_REQUESTS` (in-flight at the wrapper) and `EMB_RATE_LIMIT_RPS` (token-bucket) bound the load reaching the GPU.
|
||||
- **GPU memory**: bge-m3 fits comfortably in <4 GB; `EMB_VLLM_GPU_UTIL=0.50` is intentionally low to allow GPU sharing with reranker / other workloads.
|
||||
|
||||
---
|
||||
|
||||
## Related
|
||||
|
||||
- **`didi-brain`** consumes via `shared/embedding_client.py` for both ingestion (text → pgvector storage) and `/v1/gather` query embedding.
|
||||
- **`rerank` module** (separate, port 54200) handles cross-encoder scoring on top of dense kNN candidates from this module — they're complementary, not chained inside this service.
|
||||
- **Catalog API** uses the same endpoint for retrieval-side embeddings.
|
||||
- **pgvector** in the brain Postgres stores the resulting 1024-dim vectors (cosine distance index).
|
||||
- **Local Python use**: `from embeddings import EmbeddingClient` (works without the HTTP wrapper if you want in-process inference and have the `vllm`/`llamacpp` extras installed).
|
||||
183
ai_platform/modules/embeddings/README.md
Normal file
183
ai_platform/modules/embeddings/README.md
Normal file
|
|
@ -0,0 +1,183 @@
|
|||
# Embeddings Module
|
||||
|
||||
OpenAI-compatible embeddings API with support for vLLM and llama.cpp backends.
|
||||
|
||||
## Features
|
||||
|
||||
- **OpenAI-compatible API**: Drop-in replacement for OpenAI's `/v1/embeddings` 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/embeddings
|
||||
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-m3 \
|
||||
--host 0.0.0.0 --port 54101 \
|
||||
--task embed
|
||||
```
|
||||
|
||||
**llama.cpp (CPU):**
|
||||
```bash
|
||||
llama-server \
|
||||
--model /models/bge-m3-q4_k_m.gguf \
|
||||
--host 0.0.0.0 --port 54110 \
|
||||
--embedding
|
||||
```
|
||||
|
||||
### 4. Start API server
|
||||
|
||||
```bash
|
||||
# Set required environment variables
|
||||
export EMB_DEFAULT_BACKEND=vllm
|
||||
export EMB_ENABLE_VLLM=true
|
||||
export EMB_ENABLE_LLAMACPP=false
|
||||
export EMB_EXTERNAL_URL=http://localhost:54100
|
||||
|
||||
# Run the server
|
||||
uv run python -m embeddings.cli --port 54100
|
||||
```
|
||||
|
||||
### 5. Test the API
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:54100/v1/embeddings \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, world!",
|
||||
"model": "BAAI/bge-m3"
|
||||
}'
|
||||
```
|
||||
|
||||
## Docker Deployment
|
||||
|
||||
```bash
|
||||
cd modules/embeddings/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 embeddings import EmbeddingClient
|
||||
|
||||
# Initialize client (reads config from environment)
|
||||
client = EmbeddingClient()
|
||||
|
||||
# Generate embeddings
|
||||
response = await client.embed(
|
||||
texts=["Hello, world!", "How are you?"],
|
||||
model="BAAI/bge-m3",
|
||||
)
|
||||
|
||||
# Access embeddings
|
||||
for item in response.data:
|
||||
print(f"Index {item.index}: {len(item.embedding)} dimensions")
|
||||
|
||||
# List available models
|
||||
models = await client.list_models()
|
||||
for model in models:
|
||||
print(f"{model.id} on {model.backend}")
|
||||
```
|
||||
|
||||
## Configuration
|
||||
|
||||
All configuration is via environment variables with the `EMB_` prefix:
|
||||
|
||||
| Variable | Required | Default | Description |
|
||||
|----------|----------|---------|-------------|
|
||||
| `EMB_DEFAULT_BACKEND` | Yes | - | Default backend: `vllm` or `llamacpp` |
|
||||
| `EMB_ENABLE_VLLM` | Yes | - | Enable vLLM backend |
|
||||
| `EMB_ENABLE_LLAMACPP` | Yes | - | Enable llama.cpp backend |
|
||||
| `EMB_EXTERNAL_URL` | Yes | - | External URL for OpenAPI spec |
|
||||
| `EMB_PORT` | No | 54100 | API server port |
|
||||
| `EMB_VLLM_BASE_URL` | No | http://localhost:54101 | vLLM server URL |
|
||||
| `EMB_LLAMACPP_BASE_URL` | No | http://localhost:54110 | llama.cpp server URL |
|
||||
| `EMB_API_TOKENS` | No | - | Comma-separated API tokens |
|
||||
| `EMB_RATE_LIMIT_RPS` | No | 20.0 | Requests per second limit |
|
||||
| `EMB_MAX_CONCURRENT_REQUESTS` | No | 20 | Max concurrent requests |
|
||||
|
||||
See `.env.example` for the complete list.
|
||||
|
||||
## Port Allocation
|
||||
|
||||
Following the datacenter port schema (x41xx = Embeddings):
|
||||
|
||||
| Port | Service | Environment |
|
||||
|------|---------|-------------|
|
||||
| 14100 | Embeddings API | Production |
|
||||
| 54100 | Embeddings API | Development |
|
||||
| 14101 | vLLM Embed Server | Production |
|
||||
| 54101 | vLLM Embed Server | Development |
|
||||
| 14110 | llama.cpp Embed Server | Production |
|
||||
| 54110 | llama.cpp Embed 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/embeddings --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/embeddings/deploy/Dockerfile
Normal file
43
ai_platform/modules/embeddings/deploy/Dockerfile
Normal file
|
|
@ -0,0 +1,43 @@
|
|||
# Embeddings 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 EMB_HOST=0.0.0.0
|
||||
ENV EMB_PORT=14100
|
||||
|
||||
# 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:14100/health')" || exit 1
|
||||
|
||||
# Expose port
|
||||
EXPOSE 14100
|
||||
|
||||
# Run the server
|
||||
CMD ["python", "-m", "embeddings.cli"]
|
||||
142
ai_platform/modules/embeddings/deploy/deploy.sh
Normal file
142
ai_platform/modules/embeddings/deploy/deploy.sh
Normal file
|
|
@ -0,0 +1,142 @@
|
|||
#!/usr/bin/env bash
|
||||
#
|
||||
# Docker Compose Startup Script for Embeddings
|
||||
#
|
||||
# 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 "EMB_DEFAULT_BACKEND"
|
||||
check_required_var "EMB_ENABLE_VLLM"
|
||||
check_required_var "EMB_ENABLE_LLAMACPP"
|
||||
check_required_var "EMB_EXTERNAL_URL"
|
||||
|
||||
# Check profile-specific variables
|
||||
case $PROFILE in
|
||||
vllm)
|
||||
check_required_var "EMB_VLLM_MODEL"
|
||||
;;
|
||||
llamacpp)
|
||||
check_required_var "MODELS_DIR"
|
||||
check_required_var "EMB_LLAMACPP_MODEL"
|
||||
;;
|
||||
esac
|
||||
|
||||
cd "$SCRIPT_DIR"
|
||||
|
||||
case $ACTION in
|
||||
up)
|
||||
echo "Starting Embeddings with profile: $PROFILE"
|
||||
echo " Default backend: $EMB_DEFAULT_BACKEND"
|
||||
echo " vLLM enabled: $EMB_ENABLE_VLLM"
|
||||
echo " llama.cpp enabled: $EMB_ENABLE_LLAMACPP"
|
||||
if [[ "$PROFILE" == "vllm" ]]; then
|
||||
echo " vLLM model: ${EMB_VLLM_MODEL:-BAAI/bge-m3}"
|
||||
fi
|
||||
if [[ "$PROFILE" == "llamacpp" ]]; then
|
||||
echo " llama.cpp model: ${EMB_LLAMACPP_MODEL:-bge-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 Embeddings 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/embeddings/deploy/docker-compose.yml
Normal file
131
ai_platform/modules/embeddings/deploy/docker-compose.yml
Normal file
|
|
@ -0,0 +1,131 @@
|
|||
# Embeddings Module - Docker Compose Configuration
|
||||
#
|
||||
# Port Allocation (x41xx = Embeddings):
|
||||
# 14100 - Embeddings API (Prod)
|
||||
# 54100 - Embeddings API (Dev)
|
||||
# 14101/54101 - vLLM Embedding Server
|
||||
# 14110/54110 - llama.cpp Embedding Server
|
||||
#
|
||||
# Profiles:
|
||||
# api - API server only (uses external embedding 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:
|
||||
# ==========================================================================
|
||||
# Embeddings API Server
|
||||
# ==========================================================================
|
||||
embeddings-api:
|
||||
container_name: didiAI-embeddings-api
|
||||
image: didiai-embeddings-api
|
||||
build:
|
||||
context: ..
|
||||
dockerfile: deploy/Dockerfile
|
||||
ports:
|
||||
- "${EMB_PORT:-14100}:${EMB_PORT:-14100}"
|
||||
networks:
|
||||
- deploy_default
|
||||
environment:
|
||||
- EMB_PORT=${EMB_PORT:-14100}
|
||||
- EMB_EXTERNAL_URL=${EMB_EXTERNAL_URL}
|
||||
- EMB_DEFAULT_BACKEND=${EMB_DEFAULT_BACKEND}
|
||||
- EMB_ENABLE_VLLM=${EMB_ENABLE_VLLM}
|
||||
- EMB_ENABLE_LLAMACPP=${EMB_ENABLE_LLAMACPP}
|
||||
- EMB_VLLM_BASE_URL=http://didiAI-embeddings-vllm:14101
|
||||
- EMB_LLAMACPP_BASE_URL=http://didiAI-embeddings-llamacpp:8080
|
||||
- EMB_API_TOKENS=${EMB_API_TOKENS:-}
|
||||
- EMB_DASHBOARD_URL=${EMB_DASHBOARD_URL:-http://didiAI-dashboard:51300}
|
||||
healthcheck:
|
||||
test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:${EMB_PORT:-14100}/health')"]
|
||||
interval: 30s
|
||||
timeout: 10s
|
||||
retries: 3
|
||||
start_period: 30s
|
||||
restart: unless-stopped
|
||||
profiles:
|
||||
- api
|
||||
- vllm
|
||||
- llamacpp
|
||||
|
||||
# ==========================================================================
|
||||
# vLLM Embedding Server
|
||||
# ==========================================================================
|
||||
# Embedding model using vLLM with --task embed (v0.8.x syntax)
|
||||
vllm-embed:
|
||||
container_name: didiAI-embeddings-vllm
|
||||
image: vllm/vllm-openai:v0.8.5
|
||||
ports:
|
||||
- "${EMB_VLLM_PORT:-14101}:14101"
|
||||
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=${EMB_VLLM_GPU:-0}
|
||||
command: >
|
||||
--model ${EMB_VLLM_MODEL:-BAAI/bge-m3}
|
||||
--host 0.0.0.0
|
||||
--port 14101
|
||||
--task embed
|
||||
--trust-remote-code
|
||||
--max-model-len ${EMB_VLLM_MAX_LEN:-8192}
|
||||
--gpu-memory-utilization ${EMB_VLLM_GPU_UTIL:-0.50}
|
||||
--disable-log-requests
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
device_ids: ['${EMB_VLLM_GPU:-0}']
|
||||
capabilities: [gpu]
|
||||
healthcheck:
|
||||
test: ["CMD", "curl", "-f", "http://localhost:14101/health"]
|
||||
interval: 30s
|
||||
timeout: 10s
|
||||
retries: 10
|
||||
start_period: 300s
|
||||
restart: unless-stopped
|
||||
profiles:
|
||||
- vllm
|
||||
|
||||
# ==========================================================================
|
||||
# llama.cpp Embedding Server
|
||||
# ==========================================================================
|
||||
# Embedding model using llama.cpp with --embedding
|
||||
llamacpp-embed:
|
||||
container_name: didiAI-embeddings-llamacpp
|
||||
image: ghcr.io/ggml-org/llama.cpp:server-b4769
|
||||
ports:
|
||||
- "${EMB_LLAMACPP_PORT:-14110}:8080"
|
||||
networks:
|
||||
- deploy_default
|
||||
volumes:
|
||||
- ${MODELS_DIR:-/cai2_ds_storage/models}:/models:ro
|
||||
command: >
|
||||
--model /models/${EMB_LLAMACPP_MODEL:-bge-m3-q4_k_m.gguf}
|
||||
--host 0.0.0.0
|
||||
--port 8080
|
||||
--embedding
|
||||
--ctx-size ${EMB_LLAMACPP_CTX:-8192}
|
||||
--threads ${EMB_LLAMACPP_THREADS:-4}
|
||||
--parallel ${EMB_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/embeddings/pyproject.toml
Normal file
75
ai_platform/modules/embeddings/pyproject.toml
Normal file
|
|
@ -0,0 +1,75 @@
|
|||
[project]
|
||||
name = "embeddings"
|
||||
version = "0.1.0"
|
||||
description = "OpenAI-compatible embeddings 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 = [
|
||||
"embeddings[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]
|
||||
embeddings = "embeddings.cli:main"
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["src/embeddings"]
|
||||
|
||||
[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/embeddings/src/embeddings/__init__.py
Normal file
15
ai_platform/modules/embeddings/src/embeddings/__init__.py
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
"""Embeddings - OpenAI-compatible embeddings API with multiple backends."""
|
||||
|
||||
from embeddings.client import EmbeddingClient
|
||||
from embeddings.config import EmbeddingSettings, get_settings
|
||||
from embeddings.types import BackendType
|
||||
|
||||
__version__ = "0.1.0"
|
||||
|
||||
__all__ = [
|
||||
"BackendType",
|
||||
"EmbeddingClient",
|
||||
"EmbeddingSettings",
|
||||
"__version__",
|
||||
"get_settings",
|
||||
]
|
||||
|
|
@ -0,0 +1 @@
|
|||
"""FastAPI application for embeddings API."""
|
||||
132
ai_platform/modules/embeddings/src/embeddings/api/app.py
Normal file
132
ai_platform/modules/embeddings/src/embeddings/api/app.py
Normal file
|
|
@ -0,0 +1,132 @@
|
|||
"""FastAPI application factory."""
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
from embeddings.api.dependencies import init_concurrency_limiter
|
||||
from embeddings.api.middleware import RateLimitMiddleware, RequestIdMiddleware
|
||||
from embeddings.api.routes import embeddings, health, models
|
||||
from embeddings.client import EmbeddingClient
|
||||
from embeddings.config import SettingsCache
|
||||
from embeddings.logging import configure_logging, get_logger
|
||||
from embeddings.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 Embeddings 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_requests)
|
||||
|
||||
app.state.settings = settings
|
||||
app.state.client = EmbeddingClient(settings)
|
||||
|
||||
if hasattr(app.state, "runtime_config"):
|
||||
await app.state.runtime_config.start()
|
||||
|
||||
logger.info("Embeddings API started successfully")
|
||||
|
||||
yield
|
||||
|
||||
logger.info("Shutting down Embeddings 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="Embeddings API",
|
||||
description="OpenAI-compatible embeddings with multiple backends (vLLM, llama.cpp)",
|
||||
version="0.1.0",
|
||||
lifespan=lifespan,
|
||||
servers=[{"url": settings.external_url, "description": "Embeddings 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-embeddings-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-embeddings-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="embeddings",
|
||||
live_log_key="embeddings.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="embeddings.rate_limit.rps",
|
||||
burst_key="embeddings.rate_limit.burst",
|
||||
)
|
||||
|
||||
app.add_middleware(RequestIdMiddleware)
|
||||
|
||||
app.include_router(health.router, tags=["Health"])
|
||||
app.include_router(embeddings.router, prefix="/v1", tags=["Embeddings"])
|
||||
app.include_router(models.router, prefix="/v1", tags=["Models"])
|
||||
|
||||
return app
|
||||
|
|
@ -0,0 +1,159 @@
|
|||
"""FastAPI dependencies for embeddings."""
|
||||
|
||||
import asyncio
|
||||
import hmac
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import Header, HTTPException, Request
|
||||
|
||||
from embeddings.client import EmbeddingClient
|
||||
from embeddings.config import EmbeddingSettings
|
||||
from embeddings.logging import get_logger
|
||||
|
||||
logger = get_logger("dependencies")
|
||||
|
||||
|
||||
def get_client(request: Request) -> EmbeddingClient:
|
||||
"""Get the embedding client from application state."""
|
||||
return request.app.state.client
|
||||
|
||||
|
||||
def get_settings(request: Request) -> EmbeddingSettings:
|
||||
"""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_embedding_slot(request: Request) -> AsyncIterator[None]:
|
||||
"""FastAPI dependency that acquires an embedding slot."""
|
||||
limiter = get_concurrency_limiter()
|
||||
async with limiter.acquire():
|
||||
yield
|
||||
156
ai_platform/modules/embeddings/src/embeddings/api/middleware.py
Normal file
156
ai_platform/modules/embeddings/src/embeddings/api/middleware.py
Normal 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 embeddings.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., "embeddings.rate_limit.rps").
|
||||
burst_key: Dashboard config key (e.g., "embeddings.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 = 40,
|
||||
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)
|
||||
|
|
@ -0,0 +1 @@
|
|||
"""API route modules."""
|
||||
|
|
@ -0,0 +1,94 @@
|
|||
"""Embedding routes."""
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
|
||||
from embeddings.api.dependencies import (
|
||||
get_client,
|
||||
get_concurrency_limiter,
|
||||
verify_bearer_token,
|
||||
)
|
||||
from embeddings.client import EmbeddingClient
|
||||
from embeddings.exceptions import (
|
||||
BackendNotAvailableError,
|
||||
BackendNotEnabledError,
|
||||
EmbeddingConnectionError,
|
||||
EmbeddingRateLimitError,
|
||||
EmbeddingRequestError,
|
||||
EmbeddingTimeoutError,
|
||||
)
|
||||
from embeddings.logging import get_logger
|
||||
from embeddings.schemas import (
|
||||
EmbeddingData,
|
||||
EmbeddingRequest,
|
||||
EmbeddingResponse,
|
||||
encode_embedding_base64,
|
||||
)
|
||||
|
||||
router = APIRouter(dependencies=[Depends(verify_bearer_token)])
|
||||
logger = get_logger("routes.embeddings")
|
||||
|
||||
|
||||
@router.post("/embeddings", response_model=EmbeddingResponse)
|
||||
async def create_embeddings(
|
||||
request: EmbeddingRequest,
|
||||
client: EmbeddingClient = Depends(get_client),
|
||||
) -> EmbeddingResponse:
|
||||
"""Create embeddings for the given input.
|
||||
|
||||
This endpoint is OpenAI-compatible.
|
||||
|
||||
Args:
|
||||
request: Embedding request with input texts and model.
|
||||
client: Embedding client instance.
|
||||
|
||||
Returns:
|
||||
EmbeddingResponse: The embedding response.
|
||||
|
||||
Raises:
|
||||
HTTPException: On various error conditions.
|
||||
"""
|
||||
limiter = get_concurrency_limiter()
|
||||
|
||||
async with limiter.acquire():
|
||||
try:
|
||||
response = await client.embed(
|
||||
texts=request.input,
|
||||
model=request.model,
|
||||
backend=request.backend,
|
||||
dimensions=request.dimensions,
|
||||
)
|
||||
|
||||
# Handle encoding format
|
||||
if request.encoding_format == "base64":
|
||||
response.data = [
|
||||
EmbeddingData(
|
||||
index=d.index,
|
||||
embedding=encode_embedding_base64(d.embedding), # type: ignore[arg-type]
|
||||
)
|
||||
for d in response.data
|
||||
]
|
||||
|
||||
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 EmbeddingRateLimitError 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 EmbeddingTimeoutError as e:
|
||||
logger.warning("Timeout: %s", str(e))
|
||||
raise HTTPException(status_code=504, detail=str(e)) from None
|
||||
except EmbeddingConnectionError as e:
|
||||
logger.error("Connection error: %s", str(e))
|
||||
raise HTTPException(status_code=503, detail=str(e)) from None
|
||||
except EmbeddingRequestError as e:
|
||||
logger.error("Embedding 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
|
||||
|
|
@ -0,0 +1,69 @@
|
|||
"""Health check routes."""
|
||||
|
||||
from fastapi import APIRouter, Depends
|
||||
|
||||
from embeddings.api.dependencies import get_client
|
||||
from embeddings.client import EmbeddingClient
|
||||
from embeddings.logging import get_logger
|
||||
from embeddings.schemas import BackendHealth, HealthResponse, ReadinessResponse
|
||||
|
||||
router = APIRouter()
|
||||
logger = get_logger("routes.health")
|
||||
|
||||
|
||||
@router.get("/health", response_model=HealthResponse)
|
||||
async def health_check(
|
||||
client: EmbeddingClient = 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: EmbeddingClient = 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)
|
||||
|
|
@ -0,0 +1,63 @@
|
|||
"""Model management routes."""
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
|
||||
from embeddings.api.dependencies import get_client, verify_bearer_token
|
||||
from embeddings.client import EmbeddingClient
|
||||
from embeddings.exceptions import (
|
||||
BackendNotAvailableError,
|
||||
BackendNotEnabledError,
|
||||
ModelListError,
|
||||
)
|
||||
from embeddings.schemas import BackendListResponse, ModelListResponse
|
||||
from embeddings.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: EmbeddingClient = Depends(get_client),
|
||||
) -> ModelListResponse:
|
||||
"""List available models.
|
||||
|
||||
Args:
|
||||
backend: Optional backend filter. If not specified, returns models
|
||||
from all available backends.
|
||||
client: Embedding 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: EmbeddingClient = Depends(get_client),
|
||||
) -> BackendListResponse:
|
||||
"""List available backends.
|
||||
|
||||
Args:
|
||||
client: Embedding client instance.
|
||||
|
||||
Returns:
|
||||
BackendListResponse: List of available backend names.
|
||||
"""
|
||||
backends = [b.value for b in client.list_backends()]
|
||||
return BackendListResponse(backends=backends)
|
||||
|
|
@ -0,0 +1,9 @@
|
|||
"""Backend implementations for embeddings."""
|
||||
|
||||
from embeddings.backends.base import EmbeddingBackend
|
||||
from embeddings.backends.registry import BackendRegistry
|
||||
|
||||
__all__ = [
|
||||
"BackendRegistry",
|
||||
"EmbeddingBackend",
|
||||
]
|
||||
|
|
@ -0,0 +1,64 @@
|
|||
"""Abstract base class for embedding backends."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
from embeddings.types import EmbeddingUsage, ModelInfo
|
||||
|
||||
|
||||
class EmbeddingBackend(ABC):
|
||||
"""Abstract base class for embedding 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 embed(
|
||||
self,
|
||||
texts: list[str],
|
||||
model: str,
|
||||
dimensions: int | None = None,
|
||||
**kwargs: object,
|
||||
) -> tuple[list[list[float]], EmbeddingUsage]:
|
||||
"""Generate embeddings for texts.
|
||||
|
||||
Args:
|
||||
texts: List of texts to embed.
|
||||
model: Model identifier to use.
|
||||
dimensions: Optional desired embedding dimensions.
|
||||
**kwargs: Additional parameters passed to the backend.
|
||||
|
||||
Returns:
|
||||
Tuple of (list of embedding vectors, usage info).
|
||||
|
||||
Raises:
|
||||
EmbeddingRequestError: If the embedding 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
|
||||
|
|
@ -0,0 +1,204 @@
|
|||
"""llama.cpp backend implementation for embeddings.
|
||||
|
||||
llama.cpp supports embedding models via its OpenAI-compatible API when started
|
||||
with --embedding.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
import httpx
|
||||
|
||||
from embeddings.backends.base import EmbeddingBackend
|
||||
from embeddings.config import EmbeddingSettings
|
||||
from embeddings.exceptions import (
|
||||
EmbeddingConnectionError,
|
||||
EmbeddingRateLimitError,
|
||||
EmbeddingRequestError,
|
||||
EmbeddingTimeoutError,
|
||||
ModelListError,
|
||||
)
|
||||
from embeddings.types import EmbeddingUsage, ModelInfo
|
||||
|
||||
try:
|
||||
import openai
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
OPENAI_AVAILABLE = True
|
||||
except ImportError:
|
||||
OPENAI_AVAILABLE = False
|
||||
|
||||
logger = logging.getLogger("embeddings.backends.llamacpp")
|
||||
|
||||
|
||||
class LlamaCppEmbeddingBackend(EmbeddingBackend):
|
||||
"""llama.cpp backend using OpenAI-compatible API.
|
||||
|
||||
Requires llama.cpp server running with --embedding flag.
|
||||
|
||||
Example:
|
||||
Start llama.cpp server:
|
||||
```bash
|
||||
llama-server \
|
||||
--model /models/bge-m3-q4_k_m.gguf \
|
||||
--host 0.0.0.0 --port 54110 \
|
||||
--embedding
|
||||
```
|
||||
|
||||
Configure:
|
||||
```bash
|
||||
export EMB_ENABLE_LLAMACPP=true
|
||||
export EMB_LLAMACPP_BASE_URL=http://localhost:54110
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, settings: EmbeddingSettings) -> None:
|
||||
"""Initialize llama.cpp backend.
|
||||
|
||||
Args:
|
||||
settings: Application settings.
|
||||
|
||||
Raises:
|
||||
ImportError: If openai package is not installed.
|
||||
"""
|
||||
if not OPENAI_AVAILABLE:
|
||||
raise ImportError(
|
||||
"openai package is required for llama.cpp backend. "
|
||||
"Install with: pip install embeddings[llamacpp]"
|
||||
)
|
||||
|
||||
self._settings = settings
|
||||
|
||||
timeout = httpx.Timeout(
|
||||
connect=settings.connect_timeout,
|
||||
read=settings.request_timeout,
|
||||
write=settings.request_timeout,
|
||||
pool=settings.connect_timeout,
|
||||
)
|
||||
|
||||
self._client = AsyncOpenAI(
|
||||
base_url=f"{settings.llamacpp_base_url}/v1",
|
||||
api_key="not-needed",
|
||||
timeout=timeout,
|
||||
max_retries=0,
|
||||
)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
"""Backend identifier name."""
|
||||
return "llamacpp"
|
||||
|
||||
async def embed(
|
||||
self,
|
||||
texts: list[str],
|
||||
model: str,
|
||||
dimensions: int | None = None,
|
||||
**kwargs: object,
|
||||
) -> tuple[list[list[float]], EmbeddingUsage]:
|
||||
"""Generate embeddings via llama.cpp.
|
||||
|
||||
Args:
|
||||
texts: List of texts to embed.
|
||||
model: Model identifier (may be ignored by llama.cpp).
|
||||
dimensions: Optional desired dimensions (may not be supported).
|
||||
**kwargs: Additional parameters.
|
||||
|
||||
Returns:
|
||||
Tuple of (embeddings, usage).
|
||||
|
||||
Raises:
|
||||
EmbeddingRequestError: If the request fails.
|
||||
"""
|
||||
try:
|
||||
# llama.cpp may not support all parameters
|
||||
request_kwargs: dict[str, object] = {
|
||||
"input": texts,
|
||||
"model": model,
|
||||
}
|
||||
# Note: llama.cpp may not support dimensions parameter
|
||||
|
||||
response = await self._client.embeddings.create(**request_kwargs)
|
||||
|
||||
# Extract embeddings in order
|
||||
embeddings = [item.embedding for item in response.data]
|
||||
|
||||
# llama.cpp may not return usage, provide defaults
|
||||
usage = EmbeddingUsage(
|
||||
prompt_tokens=getattr(response.usage, "prompt_tokens", 0),
|
||||
total_tokens=getattr(response.usage, "total_tokens", 0),
|
||||
)
|
||||
|
||||
return embeddings, usage
|
||||
|
||||
except openai.RateLimitError as e:
|
||||
logger.warning("Rate limited: %s", str(e))
|
||||
raise EmbeddingRateLimitError(str(e), backend=self.name) from e
|
||||
except openai.APITimeoutError as e:
|
||||
logger.warning("Timeout: %s", str(e))
|
||||
raise EmbeddingTimeoutError(str(e), backend=self.name) from e
|
||||
except openai.APIConnectionError as e:
|
||||
logger.error("Connection error: %s", str(e))
|
||||
raise EmbeddingConnectionError(self.name, reason=str(e)) from e
|
||||
except openai.AuthenticationError as e:
|
||||
logger.error("Authentication failed: %s", str(e))
|
||||
raise EmbeddingRequestError(
|
||||
f"Authentication failed: {e}", backend=self.name, model=model
|
||||
) from e
|
||||
except openai.APIStatusError as e:
|
||||
logger.error("API error: %s", str(e))
|
||||
raise EmbeddingRequestError(str(e), backend=self.name, model=model) from e
|
||||
except Exception as e:
|
||||
logger.error("Unexpected error: %s", str(e))
|
||||
raise EmbeddingRequestError(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._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._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
|
||||
|
|
@ -0,0 +1,120 @@
|
|||
"""Backend registry for managing embedding backend instances."""
|
||||
|
||||
from embeddings.backends.base import EmbeddingBackend
|
||||
from embeddings.config import EmbeddingSettings
|
||||
from embeddings.exceptions import BackendNotAvailableError, BackendNotEnabledError
|
||||
from embeddings.types import BackendType
|
||||
|
||||
|
||||
class BackendRegistry:
|
||||
"""Registry for managing embedding 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: EmbeddingSettings) -> None:
|
||||
"""Initialize the backend registry.
|
||||
|
||||
Args:
|
||||
settings: Application settings.
|
||||
"""
|
||||
self._settings = settings
|
||||
self._backends: dict[BackendType, EmbeddingBackend] = {}
|
||||
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 embeddings.backends.vllm_backend import VLLMEmbeddingBackend
|
||||
|
||||
self._backends[BackendType.VLLM] = VLLMEmbeddingBackend(self._settings)
|
||||
except ImportError as e:
|
||||
raise BackendNotAvailableError(
|
||||
"vllm",
|
||||
f"vLLM backend is enabled but dependencies are not installed. "
|
||||
f"Install with: pip install embeddings[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 embeddings.backends.llamacpp_backend import (
|
||||
LlamaCppEmbeddingBackend,
|
||||
)
|
||||
|
||||
self._backends[BackendType.LLAMACPP] = LlamaCppEmbeddingBackend(
|
||||
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 embeddings[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: "
|
||||
"EMB_ENABLE_VLLM=true or EMB_ENABLE_LLAMACPP=true",
|
||||
)
|
||||
|
||||
def get(self, backend_type: BackendType | str | None = None) -> EmbeddingBackend:
|
||||
"""Get a backend instance.
|
||||
|
||||
Args:
|
||||
backend_type: Backend to retrieve. If None, uses default from settings.
|
||||
|
||||
Returns:
|
||||
EmbeddingBackend: 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:
|
||||
# Check if it's a valid backend that's just not enabled
|
||||
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.
|
||||
|
||||
Returns:
|
||||
list[BackendType]: List of available backend types.
|
||||
"""
|
||||
return list(self._backends.keys())
|
||||
|
||||
def is_available(self, backend_type: BackendType | str) -> bool:
|
||||
"""Check if a backend is available.
|
||||
|
||||
Args:
|
||||
backend_type: Backend to check.
|
||||
|
||||
Returns:
|
||||
bool: True if the backend is available.
|
||||
"""
|
||||
if isinstance(backend_type, str):
|
||||
backend_type = BackendType(backend_type)
|
||||
return backend_type in self._backends
|
||||
|
|
@ -0,0 +1,202 @@
|
|||
"""vLLM backend implementation for embeddings.
|
||||
|
||||
vLLM supports embedding models via the OpenAI-compatible API when started
|
||||
with --task embed.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
import httpx
|
||||
|
||||
from embeddings.backends.base import EmbeddingBackend
|
||||
from embeddings.config import EmbeddingSettings
|
||||
from embeddings.exceptions import (
|
||||
EmbeddingConnectionError,
|
||||
EmbeddingRateLimitError,
|
||||
EmbeddingRequestError,
|
||||
EmbeddingTimeoutError,
|
||||
ModelListError,
|
||||
)
|
||||
from embeddings.types import EmbeddingUsage, ModelInfo
|
||||
|
||||
try:
|
||||
import openai
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
OPENAI_AVAILABLE = True
|
||||
except ImportError:
|
||||
OPENAI_AVAILABLE = False
|
||||
|
||||
logger = logging.getLogger("embeddings.backends.vllm")
|
||||
|
||||
|
||||
class VLLMEmbeddingBackend(EmbeddingBackend):
|
||||
"""vLLM backend using OpenAI-compatible API.
|
||||
|
||||
Requires the vLLM server to be running with --task embed.
|
||||
|
||||
Example:
|
||||
Start vLLM server:
|
||||
```bash
|
||||
python -m vllm.entrypoints.openai.api_server \
|
||||
--model BAAI/bge-m3 \
|
||||
--host 0.0.0.0 --port 54101 \
|
||||
--task embed
|
||||
```
|
||||
|
||||
Configure:
|
||||
```bash
|
||||
export EMB_ENABLE_VLLM=true
|
||||
export EMB_VLLM_BASE_URL=http://localhost:54101
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, settings: EmbeddingSettings) -> None:
|
||||
"""Initialize vLLM backend.
|
||||
|
||||
Args:
|
||||
settings: Application settings.
|
||||
|
||||
Raises:
|
||||
ImportError: If openai package is not installed.
|
||||
"""
|
||||
if not OPENAI_AVAILABLE:
|
||||
raise ImportError(
|
||||
"openai package is required for vLLM backend. "
|
||||
"Install with: pip install embeddings[vllm]"
|
||||
)
|
||||
|
||||
self._settings = settings
|
||||
|
||||
timeout = httpx.Timeout(
|
||||
connect=settings.connect_timeout,
|
||||
read=settings.request_timeout,
|
||||
write=settings.request_timeout,
|
||||
pool=settings.connect_timeout,
|
||||
)
|
||||
|
||||
self._client = AsyncOpenAI(
|
||||
base_url=f"{settings.vllm_base_url}/v1",
|
||||
api_key=settings.vllm_api_key or "not-needed",
|
||||
timeout=timeout,
|
||||
max_retries=0,
|
||||
)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
"""Backend identifier name."""
|
||||
return "vllm"
|
||||
|
||||
async def embed(
|
||||
self,
|
||||
texts: list[str],
|
||||
model: str,
|
||||
dimensions: int | None = None,
|
||||
**kwargs: object,
|
||||
) -> tuple[list[list[float]], EmbeddingUsage]:
|
||||
"""Generate embeddings via vLLM.
|
||||
|
||||
Args:
|
||||
texts: List of texts to embed.
|
||||
model: Model identifier.
|
||||
dimensions: Optional desired dimensions.
|
||||
**kwargs: Additional parameters.
|
||||
|
||||
Returns:
|
||||
Tuple of (embeddings, usage).
|
||||
|
||||
Raises:
|
||||
EmbeddingRequestError: If the request fails.
|
||||
"""
|
||||
try:
|
||||
# Build request parameters
|
||||
request_kwargs: dict[str, object] = {
|
||||
"input": texts,
|
||||
"model": model,
|
||||
}
|
||||
if dimensions is not None:
|
||||
request_kwargs["dimensions"] = dimensions
|
||||
|
||||
response = await self._client.embeddings.create(**request_kwargs)
|
||||
|
||||
# Extract embeddings in order
|
||||
embeddings = [item.embedding for item in response.data]
|
||||
|
||||
usage = EmbeddingUsage(
|
||||
prompt_tokens=response.usage.prompt_tokens,
|
||||
total_tokens=response.usage.total_tokens,
|
||||
)
|
||||
|
||||
return embeddings, usage
|
||||
|
||||
except openai.RateLimitError as e:
|
||||
logger.warning("Rate limited: %s", str(e))
|
||||
raise EmbeddingRateLimitError(str(e), backend=self.name) from e
|
||||
except openai.APITimeoutError as e:
|
||||
logger.warning("Timeout: %s", str(e))
|
||||
raise EmbeddingTimeoutError(str(e), backend=self.name) from e
|
||||
except openai.APIConnectionError as e:
|
||||
logger.error("Connection error: %s", str(e))
|
||||
raise EmbeddingConnectionError(self.name, reason=str(e)) from e
|
||||
except openai.AuthenticationError as e:
|
||||
logger.error("Authentication failed: %s", str(e))
|
||||
raise EmbeddingRequestError(
|
||||
f"Authentication failed: {e}", backend=self.name, model=model
|
||||
) from e
|
||||
except openai.APIStatusError as e:
|
||||
logger.error("API error: %s", str(e))
|
||||
raise EmbeddingRequestError(str(e), backend=self.name, model=model) from e
|
||||
except Exception as e:
|
||||
logger.error("Unexpected error: %s", str(e))
|
||||
raise EmbeddingRequestError(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._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._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/embeddings/src/embeddings/cli.py
Normal file
123
ai_platform/modules/embeddings/src/embeddings/cli.py
Normal file
|
|
@ -0,0 +1,123 @@
|
|||
"""CLI entry point for the embeddings server."""
|
||||
|
||||
import argparse
|
||||
import signal
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
import uvicorn
|
||||
|
||||
from embeddings.config import SettingsCache
|
||||
from embeddings.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 embeddings server with graceful shutdown support."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Embeddings API Server",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--host",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Host to bind to (default: from EMB_HOST env or 0.0.0.0)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--port",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Port to bind to (default: from EMB_PORT env or 54100)",
|
||||
)
|
||||
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 Embeddings API Server on %s:%d", host, port)
|
||||
|
||||
if args.reload:
|
||||
logger.info("Development mode - auto-reload enabled")
|
||||
uvicorn.run(
|
||||
"embeddings.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(
|
||||
"embeddings.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()
|
||||
157
ai_platform/modules/embeddings/src/embeddings/client.py
Normal file
157
ai_platform/modules/embeddings/src/embeddings/client.py
Normal file
|
|
@ -0,0 +1,157 @@
|
|||
"""High-level embedding client for unified inference."""
|
||||
|
||||
from embeddings.backends import BackendRegistry
|
||||
from embeddings.config import EmbeddingSettings, SettingsCache
|
||||
from embeddings.schemas import EmbeddingData, EmbeddingResponse
|
||||
from embeddings.types import BackendType, ModelInfo
|
||||
|
||||
|
||||
class EmbeddingClient:
|
||||
"""High-level client for embedding inference.
|
||||
|
||||
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 embeddings import EmbeddingClient
|
||||
|
||||
client = EmbeddingClient()
|
||||
|
||||
# Single text embedding
|
||||
response = await client.embed(
|
||||
texts=["Hello, world!"],
|
||||
model="BAAI/bge-m3",
|
||||
)
|
||||
print(response.data[0].embedding[:5])
|
||||
|
||||
# Multiple texts
|
||||
response = await client.embed(
|
||||
texts=["Hello", "World"],
|
||||
model="BAAI/bge-m3",
|
||||
)
|
||||
|
||||
# Override backend per-request
|
||||
response = await client.embed(
|
||||
texts=["Hello"],
|
||||
model="bge-m3-q4_k_m.gguf",
|
||||
backend=BackendType.LLAMACPP,
|
||||
)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, settings: EmbeddingSettings | None = None) -> None:
|
||||
"""Initialize the embedding 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.
|
||||
|
||||
Returns:
|
||||
BackendRegistry: The backend registry instance.
|
||||
"""
|
||||
return self._registry
|
||||
|
||||
async def embed(
|
||||
self,
|
||||
texts: list[str] | str,
|
||||
model: str,
|
||||
backend: BackendType | str | None = None,
|
||||
dimensions: int | None = None,
|
||||
**kwargs: object,
|
||||
) -> EmbeddingResponse:
|
||||
"""Generate embeddings for texts.
|
||||
|
||||
Args:
|
||||
texts: Text or list of texts to embed.
|
||||
model: Model identifier.
|
||||
backend: Backend to use. Uses default if not specified.
|
||||
dimensions: Optional desired embedding dimensions.
|
||||
**kwargs: Additional parameters.
|
||||
|
||||
Returns:
|
||||
EmbeddingResponse: The embedding response.
|
||||
|
||||
Raises:
|
||||
BackendNotAvailableError: If the requested backend is not available.
|
||||
EmbeddingRequestError: If the embedding request fails.
|
||||
"""
|
||||
# Normalize texts to list
|
||||
if isinstance(texts, str):
|
||||
texts = [texts]
|
||||
|
||||
backend_instance = self._registry.get(backend)
|
||||
|
||||
embeddings, usage = await backend_instance.embed(
|
||||
texts=texts,
|
||||
model=model,
|
||||
dimensions=dimensions,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Build response
|
||||
data = [
|
||||
EmbeddingData(
|
||||
index=i,
|
||||
embedding=embedding,
|
||||
)
|
||||
for i, embedding in enumerate(embeddings)
|
||||
]
|
||||
|
||||
return EmbeddingResponse(
|
||||
data=data,
|
||||
model=model,
|
||||
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.
|
||||
|
||||
Returns:
|
||||
list[BackendType]: List of available backend types.
|
||||
"""
|
||||
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
|
||||
212
ai_platform/modules/embeddings/src/embeddings/config.py
Normal file
212
ai_platform/modules/embeddings/src/embeddings/config.py
Normal file
|
|
@ -0,0 +1,212 @@
|
|||
"""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 EmbeddingSettings(BaseSettings):
|
||||
"""Embeddings module configuration.
|
||||
|
||||
All settings can be configured via environment variables with the EMB_ prefix.
|
||||
|
||||
Required environment variables (no defaults - fail-fast):
|
||||
EMB_DEFAULT_BACKEND: Backend to use (vllm, llamacpp)
|
||||
EMB_ENABLE_VLLM: Enable vLLM backend (true/false)
|
||||
EMB_ENABLE_LLAMACPP: Enable llama.cpp backend (true/false)
|
||||
EMB_EXTERNAL_URL: External URL for OpenAPI spec
|
||||
|
||||
Example:
|
||||
EMB_DEFAULT_BACKEND=vllm
|
||||
EMB_ENABLE_VLLM=true
|
||||
EMB_ENABLE_LLAMACPP=false
|
||||
EMB_EXTERNAL_URL=http://localhost:54100
|
||||
"""
|
||||
|
||||
model_config = SettingsConfigDict(
|
||||
env_prefix="EMB_",
|
||||
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 embeddings (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=54100,
|
||||
description="Port to bind the server to (Dev: 54100, Prod: 14100)",
|
||||
)
|
||||
|
||||
# vLLM settings
|
||||
vllm_base_url: str = Field(
|
||||
default="http://localhost:54101",
|
||||
description="Base URL for vLLM embedding 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:54110",
|
||||
description="Base URL for llama.cpp embedding server",
|
||||
)
|
||||
|
||||
# ==========================================================================
|
||||
# Timeout settings
|
||||
# ==========================================================================
|
||||
|
||||
request_timeout: float = Field(
|
||||
default=120.0,
|
||||
ge=1.0,
|
||||
description="Default timeout in seconds for embedding 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=40,
|
||||
ge=1,
|
||||
description="Rate limit: burst capacity",
|
||||
)
|
||||
max_concurrent_requests: int = Field(
|
||||
default=20,
|
||||
ge=1,
|
||||
description="Maximum concurrent embedding 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: EmbeddingSettings | None = None
|
||||
_lock: threading.Lock = threading.Lock()
|
||||
|
||||
@classmethod
|
||||
def get(cls) -> EmbeddingSettings:
|
||||
"""Get or create the settings instance.
|
||||
|
||||
Returns:
|
||||
EmbeddingSettings: Application settings loaded from environment.
|
||||
"""
|
||||
with cls._lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = EmbeddingSettings()
|
||||
return cls._instance
|
||||
|
||||
@classmethod
|
||||
def clear(cls) -> None:
|
||||
"""Clear the cached settings instance."""
|
||||
with cls._lock:
|
||||
cls._instance = None
|
||||
|
||||
@classmethod
|
||||
def set(cls, settings: EmbeddingSettings) -> None:
|
||||
"""Set a specific settings instance.
|
||||
|
||||
Args:
|
||||
settings: Settings instance to use.
|
||||
"""
|
||||
with cls._lock:
|
||||
cls._instance = settings
|
||||
|
||||
|
||||
def get_settings() -> EmbeddingSettings:
|
||||
"""Get cached settings instance.
|
||||
|
||||
Returns:
|
||||
EmbeddingSettings: Application settings loaded from environment.
|
||||
"""
|
||||
return SettingsCache.get()
|
||||
124
ai_platform/modules/embeddings/src/embeddings/exceptions.py
Normal file
124
ai_platform/modules/embeddings/src/embeddings/exceptions.py
Normal file
|
|
@ -0,0 +1,124 @@
|
|||
"""Custom exceptions for embeddings module."""
|
||||
|
||||
|
||||
class EmbeddingError(Exception):
|
||||
"""Base exception for embedding errors."""
|
||||
|
||||
|
||||
class AuthenticationError(EmbeddingError):
|
||||
"""Raised when authentication fails."""
|
||||
|
||||
def __init__(self, message: str = "Authentication required") -> None:
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class BackendNotAvailableError(EmbeddingError):
|
||||
"""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(EmbeddingError):
|
||||
"""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 EMB_ENABLE_{backend.upper()}=true to enable it."
|
||||
)
|
||||
|
||||
|
||||
class ModelNotFoundError(EmbeddingError):
|
||||
"""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 EmbeddingRequestError(EmbeddingError):
|
||||
"""Raised when an embedding 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(EmbeddingError):
|
||||
"""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 EmbeddingRateLimitError(EmbeddingError):
|
||||
"""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 EmbeddingTimeoutError(EmbeddingError):
|
||||
"""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 EmbeddingConnectionError(EmbeddingError):
|
||||
"""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)
|
||||
98
ai_platform/modules/embeddings/src/embeddings/logging.py
Normal file
98
ai_platform/modules/embeddings/src/embeddings/logging.py
Normal file
|
|
@ -0,0 +1,98 @@
|
|||
"""Structured logging with request_id context propagation."""
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from contextvars import ContextVar
|
||||
from typing import Any
|
||||
|
||||
# Context variable for request ID propagation
|
||||
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.
|
||||
|
||||
Args:
|
||||
level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL).
|
||||
json_format: If True, use JSON formatting; otherwise use text.
|
||||
"""
|
||||
root_logger = logging.getLogger("embeddings")
|
||||
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 embeddings prefix.
|
||||
|
||||
Args:
|
||||
name: Logger name (will be prefixed with 'embeddings.').
|
||||
|
||||
Returns:
|
||||
logging.Logger: Configured logger instance.
|
||||
"""
|
||||
if name.startswith("embeddings."):
|
||||
return logging.getLogger(name)
|
||||
return logging.getLogger(f"embeddings.{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/embeddings/src/embeddings/runtime_config.py
Normal file
157
ai_platform/modules/embeddings/src/embeddings/runtime_config.py
Normal 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)
|
||||
138
ai_platform/modules/embeddings/src/embeddings/schemas.py
Normal file
138
ai_platform/modules/embeddings/src/embeddings/schemas.py
Normal file
|
|
@ -0,0 +1,138 @@
|
|||
"""Pydantic V2 API request/response schemas."""
|
||||
|
||||
import base64
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
|
||||
from embeddings.types import BackendType, EmbeddingData, EmbeddingUsage, ModelInfo
|
||||
|
||||
# =============================================================================
|
||||
# Embedding Schemas
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class EmbeddingRequest(BaseModel):
|
||||
"""Request schema for embeddings (OpenAI-compatible)."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
input: list[str] | str = Field(
|
||||
description="Text(s) to embed. Can be a single string or list of strings.",
|
||||
)
|
||||
model: str = Field(
|
||||
description="Model to use for embedding",
|
||||
)
|
||||
encoding_format: Literal["float", "base64"] = Field(
|
||||
default="float",
|
||||
description="Format for embedding values: 'float' (default) or 'base64'",
|
||||
)
|
||||
dimensions: int | None = Field(
|
||||
default=None,
|
||||
ge=1,
|
||||
description="Desired embedding dimensions (if model supports)",
|
||||
)
|
||||
backend: BackendType | None = Field(
|
||||
default=None,
|
||||
description="Backend to use (overrides default)",
|
||||
)
|
||||
|
||||
@field_validator("input", mode="before")
|
||||
@classmethod
|
||||
def ensure_list(cls, v: list[str] | str) -> list[str]:
|
||||
"""Convert single string to list."""
|
||||
if isinstance(v, str):
|
||||
return [v]
|
||||
return v
|
||||
|
||||
@field_validator("input")
|
||||
@classmethod
|
||||
def validate_input(cls, v: list[str]) -> list[str]:
|
||||
"""Validate input is not empty."""
|
||||
if not v:
|
||||
raise ValueError("input cannot be empty")
|
||||
if any(not text.strip() for text in v):
|
||||
raise ValueError("input cannot contain empty strings")
|
||||
return v
|
||||
|
||||
|
||||
class EmbeddingResponse(BaseModel):
|
||||
"""Response schema for embeddings (OpenAI-compatible)."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
object: str = Field(default="list")
|
||||
data: list[EmbeddingData] = Field(description="List of embedding results")
|
||||
model: str = Field(description="Model used for embedding")
|
||||
usage: EmbeddingUsage = 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")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Utility Functions
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def encode_embedding_base64(embedding: list[float]) -> str:
|
||||
"""Encode embedding as base64 string.
|
||||
|
||||
Args:
|
||||
embedding: List of float values.
|
||||
|
||||
Returns:
|
||||
Base64-encoded string of the embedding.
|
||||
"""
|
||||
import struct
|
||||
|
||||
# Pack as little-endian floats
|
||||
packed = struct.pack(f"<{len(embedding)}f", *embedding)
|
||||
return base64.b64encode(packed).decode("ascii")
|
||||
45
ai_platform/modules/embeddings/src/embeddings/types.py
Normal file
45
ai_platform/modules/embeddings/src/embeddings/types.py
Normal file
|
|
@ -0,0 +1,45 @@
|
|||
"""Core types and enums for embeddings."""
|
||||
|
||||
from enum import Enum
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
class BackendType(str, Enum):
|
||||
"""Supported embedding backends."""
|
||||
|
||||
VLLM = "vllm"
|
||||
LLAMACPP = "llamacpp"
|
||||
|
||||
|
||||
class EmbeddingUsage(BaseModel):
|
||||
"""Token usage information for embeddings."""
|
||||
|
||||
prompt_tokens: int = Field(default=0)
|
||||
total_tokens: int = Field(default=0)
|
||||
|
||||
|
||||
class EmbeddingData(BaseModel):
|
||||
"""A single embedding result."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
object: str = Field(default="embedding")
|
||||
index: int = Field(description="Index of the input text")
|
||||
embedding: list[float] = Field(description="The embedding vector")
|
||||
|
||||
|
||||
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")
|
||||
dimensions: int | None = Field(
|
||||
default=None,
|
||||
description="Embedding dimensions",
|
||||
)
|
||||
max_input_tokens: int | None = Field(
|
||||
default=None,
|
||||
description="Maximum input tokens",
|
||||
)
|
||||
1
ai_platform/modules/embeddings/tests/__init__.py
Normal file
1
ai_platform/modules/embeddings/tests/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
"""Tests for embeddings module."""
|
||||
82
ai_platform/modules/embeddings/tests/conftest.py
Normal file
82
ai_platform/modules/embeddings/tests/conftest.py
Normal file
|
|
@ -0,0 +1,82 @@
|
|||
"""Pytest fixtures for embeddings tests."""
|
||||
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from embeddings.api.app import create_app
|
||||
from embeddings.api.dependencies import init_concurrency_limiter
|
||||
from embeddings.client import EmbeddingClient
|
||||
from embeddings.config import EmbeddingSettings, SettingsCache
|
||||
from embeddings.schemas import EmbeddingResponse
|
||||
from embeddings.types import BackendType, EmbeddingData, EmbeddingUsage
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clear_settings_cache() -> None:
|
||||
"""Clear settings cache before each test."""
|
||||
SettingsCache.clear()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_settings() -> EmbeddingSettings:
|
||||
"""Create test settings with mocked values."""
|
||||
return EmbeddingSettings(
|
||||
default_backend="vllm",
|
||||
enable_vllm=True,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:54100",
|
||||
host="127.0.0.1",
|
||||
port=54100,
|
||||
api_tokens=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_embedding_response() -> EmbeddingResponse:
|
||||
"""Create a mock embedding response."""
|
||||
return EmbeddingResponse(
|
||||
data=[
|
||||
EmbeddingData(
|
||||
index=0,
|
||||
embedding=[0.1, 0.2, 0.3, 0.4, 0.5],
|
||||
)
|
||||
],
|
||||
model="test-model",
|
||||
usage=EmbeddingUsage(prompt_tokens=5, total_tokens=5),
|
||||
backend="vllm",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client_with_mock_backend(
|
||||
test_settings: EmbeddingSettings, mock_embedding_response: EmbeddingResponse
|
||||
) -> EmbeddingClient:
|
||||
"""Create an EmbeddingClient with mocked backend."""
|
||||
SettingsCache.set(test_settings)
|
||||
|
||||
client = EmbeddingClient(settings=test_settings)
|
||||
|
||||
# Mock the backend's embed method
|
||||
backend = client.registry.get(BackendType.VLLM)
|
||||
embeddings = [d.embedding for d in mock_embedding_response.data]
|
||||
usage = mock_embedding_response.usage
|
||||
backend.embed = AsyncMock(return_value=(embeddings, usage))
|
||||
|
||||
return client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app_client(test_settings: EmbeddingSettings) -> TestClient:
|
||||
"""Create a FastAPI TestClient with test settings."""
|
||||
SettingsCache.set(test_settings)
|
||||
|
||||
init_concurrency_limiter(test_settings.max_concurrent_requests)
|
||||
|
||||
app = create_app()
|
||||
|
||||
app.state.settings = test_settings
|
||||
app.state.client = EmbeddingClient(settings=test_settings)
|
||||
|
||||
return TestClient(app)
|
||||
179
ai_platform/modules/embeddings/tests/test_config.py
Normal file
179
ai_platform/modules/embeddings/tests/test_config.py
Normal file
|
|
@ -0,0 +1,179 @@
|
|||
"""Tests for configuration module."""
|
||||
|
||||
import os
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from embeddings.config import EmbeddingSettings, SettingsCache
|
||||
|
||||
|
||||
class TestEmbeddingSettings:
|
||||
"""Tests for EmbeddingSettings."""
|
||||
|
||||
def test_required_fields(self) -> None:
|
||||
"""Test that required fields must be provided."""
|
||||
with pytest.raises(ValidationError):
|
||||
EmbeddingSettings()
|
||||
|
||||
def test_with_required_fields(self) -> None:
|
||||
"""Test settings with all required fields provided."""
|
||||
settings = EmbeddingSettings(
|
||||
default_backend="vllm",
|
||||
enable_vllm=True,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:54100",
|
||||
)
|
||||
|
||||
assert settings.default_backend == "vllm"
|
||||
assert settings.enable_vllm is True
|
||||
assert settings.enable_llamacpp is False
|
||||
assert settings.external_url == "http://localhost:54100"
|
||||
|
||||
def test_optional_defaults(self) -> None:
|
||||
"""Test optional fields have sensible defaults."""
|
||||
settings = EmbeddingSettings(
|
||||
default_backend="vllm",
|
||||
enable_vllm=True,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:54100",
|
||||
)
|
||||
|
||||
assert settings.host == "0.0.0.0"
|
||||
assert settings.port == 54100
|
||||
assert settings.request_timeout == 120.0
|
||||
assert settings.rate_limit_rps == 20.0
|
||||
assert settings.max_concurrent_requests == 20
|
||||
|
||||
def test_env_override(self) -> None:
|
||||
"""Test environment variable overrides."""
|
||||
env_vars = {
|
||||
"EMB_DEFAULT_BACKEND": "llamacpp",
|
||||
"EMB_ENABLE_VLLM": "false",
|
||||
"EMB_ENABLE_LLAMACPP": "true",
|
||||
"EMB_EXTERNAL_URL": "http://example.com",
|
||||
"EMB_PORT": "14100",
|
||||
}
|
||||
|
||||
with patch.dict(os.environ, env_vars, clear=False):
|
||||
settings = EmbeddingSettings()
|
||||
|
||||
assert settings.default_backend == "llamacpp"
|
||||
assert settings.enable_vllm is False
|
||||
assert settings.enable_llamacpp is True
|
||||
assert settings.port == 14100
|
||||
|
||||
def test_vllm_url_default(self) -> None:
|
||||
"""Test vLLM URL default value."""
|
||||
settings = EmbeddingSettings(
|
||||
default_backend="vllm",
|
||||
enable_vllm=True,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:54100",
|
||||
)
|
||||
assert settings.vllm_base_url == "http://localhost:54101"
|
||||
|
||||
def test_llamacpp_url_default(self) -> None:
|
||||
"""Test llama.cpp URL default value."""
|
||||
settings = EmbeddingSettings(
|
||||
default_backend="vllm",
|
||||
enable_vllm=True,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:54100",
|
||||
)
|
||||
assert settings.llamacpp_base_url == "http://localhost:54110"
|
||||
|
||||
def test_api_tokens_parsing(self) -> None:
|
||||
"""Test API tokens are parsed from comma-separated string."""
|
||||
settings = EmbeddingSettings(
|
||||
default_backend="vllm",
|
||||
enable_vllm=True,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:54100",
|
||||
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 = EmbeddingSettings(
|
||||
default_backend="vllm",
|
||||
enable_vllm=True,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:54100",
|
||||
)
|
||||
|
||||
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 = EmbeddingSettings(
|
||||
default_backend="vllm",
|
||||
enable_vllm=True,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:54100",
|
||||
)
|
||||
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 = {
|
||||
"EMB_DEFAULT_BACKEND": "vllm",
|
||||
"EMB_ENABLE_VLLM": "true",
|
||||
"EMB_ENABLE_LLAMACPP": "false",
|
||||
"EMB_EXTERNAL_URL": "http://localhost:54100",
|
||||
}
|
||||
|
||||
with patch.dict(os.environ, env_vars, clear=False):
|
||||
settings = SettingsCache.get()
|
||||
assert isinstance(settings, EmbeddingSettings)
|
||||
assert settings.default_backend == "vllm"
|
||||
|
||||
def test_cached_returns_same_instance(self) -> None:
|
||||
"""Test that SettingsCache returns cached instance."""
|
||||
SettingsCache.clear()
|
||||
|
||||
settings = EmbeddingSettings(
|
||||
default_backend="vllm",
|
||||
enable_vllm=True,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:54100",
|
||||
)
|
||||
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 = EmbeddingSettings(
|
||||
default_backend="vllm",
|
||||
enable_vllm=True,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:54100",
|
||||
)
|
||||
SettingsCache.set(settings)
|
||||
|
||||
SettingsCache.clear()
|
||||
|
||||
assert SettingsCache._instance is None
|
||||
124
ai_platform/modules/embeddings/tests/test_schemas.py
Normal file
124
ai_platform/modules/embeddings/tests/test_schemas.py
Normal file
|
|
@ -0,0 +1,124 @@
|
|||
"""Tests for schemas module."""
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from embeddings.schemas import (
|
||||
EmbeddingRequest,
|
||||
EmbeddingResponse,
|
||||
encode_embedding_base64,
|
||||
)
|
||||
from embeddings.types import BackendType, EmbeddingData, EmbeddingUsage
|
||||
|
||||
|
||||
class TestEmbeddingRequest:
|
||||
"""Tests for EmbeddingRequest schema."""
|
||||
|
||||
def test_single_string_input(self) -> None:
|
||||
"""Test that single string input is converted to list."""
|
||||
request = EmbeddingRequest(input="hello", model="test-model")
|
||||
assert request.input == ["hello"]
|
||||
|
||||
def test_list_input(self) -> None:
|
||||
"""Test list input is preserved."""
|
||||
request = EmbeddingRequest(input=["hello", "world"], model="test-model")
|
||||
assert request.input == ["hello", "world"]
|
||||
|
||||
def test_empty_input_rejected(self) -> None:
|
||||
"""Test that empty input is rejected."""
|
||||
with pytest.raises(ValidationError):
|
||||
EmbeddingRequest(input=[], model="test-model")
|
||||
|
||||
def test_empty_string_in_list_rejected(self) -> None:
|
||||
"""Test that empty strings in list are rejected."""
|
||||
with pytest.raises(ValidationError):
|
||||
EmbeddingRequest(input=["hello", ""], model="test-model")
|
||||
|
||||
def test_default_encoding_format(self) -> None:
|
||||
"""Test default encoding format is float."""
|
||||
request = EmbeddingRequest(input="hello", model="test-model")
|
||||
assert request.encoding_format == "float"
|
||||
|
||||
def test_base64_encoding_format(self) -> None:
|
||||
"""Test base64 encoding format."""
|
||||
request = EmbeddingRequest(
|
||||
input="hello", model="test-model", encoding_format="base64"
|
||||
)
|
||||
assert request.encoding_format == "base64"
|
||||
|
||||
def test_backend_override(self) -> None:
|
||||
"""Test backend override."""
|
||||
request = EmbeddingRequest(
|
||||
input="hello", model="test-model", backend=BackendType.LLAMACPP
|
||||
)
|
||||
assert request.backend == BackendType.LLAMACPP
|
||||
|
||||
def test_dimensions_parameter(self) -> None:
|
||||
"""Test dimensions parameter."""
|
||||
request = EmbeddingRequest(input="hello", model="test-model", dimensions=256)
|
||||
assert request.dimensions == 256
|
||||
|
||||
def test_invalid_dimensions(self) -> None:
|
||||
"""Test that invalid dimensions are rejected."""
|
||||
with pytest.raises(ValidationError):
|
||||
EmbeddingRequest(input="hello", model="test-model", dimensions=0)
|
||||
|
||||
|
||||
class TestEmbeddingResponse:
|
||||
"""Tests for EmbeddingResponse schema."""
|
||||
|
||||
def test_creation(self) -> None:
|
||||
"""Test creating response."""
|
||||
response = EmbeddingResponse(
|
||||
data=[EmbeddingData(index=0, embedding=[0.1, 0.2, 0.3])],
|
||||
model="test-model",
|
||||
usage=EmbeddingUsage(prompt_tokens=5, total_tokens=5),
|
||||
backend="vllm",
|
||||
)
|
||||
assert response.object == "list"
|
||||
assert len(response.data) == 1
|
||||
assert response.model == "test-model"
|
||||
assert response.backend == "vllm"
|
||||
|
||||
def test_multiple_embeddings(self) -> None:
|
||||
"""Test response with multiple embeddings."""
|
||||
response = EmbeddingResponse(
|
||||
data=[
|
||||
EmbeddingData(index=0, embedding=[0.1, 0.2]),
|
||||
EmbeddingData(index=1, embedding=[0.3, 0.4]),
|
||||
],
|
||||
model="test-model",
|
||||
usage=EmbeddingUsage(prompt_tokens=10, total_tokens=10),
|
||||
backend="vllm",
|
||||
)
|
||||
assert len(response.data) == 2
|
||||
assert response.data[0].index == 0
|
||||
assert response.data[1].index == 1
|
||||
|
||||
|
||||
class TestEncodeEmbeddingBase64:
|
||||
"""Tests for base64 encoding function."""
|
||||
|
||||
def test_encode_simple(self) -> None:
|
||||
"""Test encoding simple embedding."""
|
||||
embedding = [1.0, 2.0, 3.0]
|
||||
encoded = encode_embedding_base64(embedding)
|
||||
assert isinstance(encoded, str)
|
||||
# Should be base64 encoded
|
||||
assert len(encoded) > 0
|
||||
|
||||
def test_encode_decode_roundtrip(self) -> None:
|
||||
"""Test that encoding can be reversed."""
|
||||
import base64
|
||||
import struct
|
||||
|
||||
embedding = [0.1, 0.2, 0.3, 0.4, 0.5]
|
||||
encoded = encode_embedding_base64(embedding)
|
||||
|
||||
# Decode
|
||||
decoded_bytes = base64.b64decode(encoded)
|
||||
decoded = list(struct.unpack(f"<{len(embedding)}f", decoded_bytes))
|
||||
|
||||
# Compare with tolerance for float precision
|
||||
for original, decoded_val in zip(embedding, decoded):
|
||||
assert abs(original - decoded_val) < 1e-6
|
||||
98
ai_platform/modules/embeddings/tests/test_types.py
Normal file
98
ai_platform/modules/embeddings/tests/test_types.py
Normal file
|
|
@ -0,0 +1,98 @@
|
|||
"""Tests for types module."""
|
||||
|
||||
import pytest
|
||||
|
||||
from embeddings.types import BackendType, EmbeddingData, EmbeddingUsage, ModelInfo
|
||||
|
||||
|
||||
class TestBackendType:
|
||||
"""Tests for BackendType enum."""
|
||||
|
||||
def test_vllm_value(self) -> None:
|
||||
"""Test vLLM backend type value."""
|
||||
assert BackendType.VLLM.value == "vllm"
|
||||
|
||||
def test_llamacpp_value(self) -> None:
|
||||
"""Test llama.cpp backend type value."""
|
||||
assert BackendType.LLAMACPP.value == "llamacpp"
|
||||
|
||||
def test_from_string(self) -> None:
|
||||
"""Test creating backend type from string."""
|
||||
assert BackendType("vllm") == BackendType.VLLM
|
||||
assert BackendType("llamacpp") == BackendType.LLAMACPP
|
||||
|
||||
def test_invalid_backend(self) -> None:
|
||||
"""Test that invalid backend raises ValueError."""
|
||||
with pytest.raises(ValueError):
|
||||
BackendType("invalid")
|
||||
|
||||
|
||||
class TestEmbeddingUsage:
|
||||
"""Tests for EmbeddingUsage model."""
|
||||
|
||||
def test_default_values(self) -> None:
|
||||
"""Test default values."""
|
||||
usage = EmbeddingUsage()
|
||||
assert usage.prompt_tokens == 0
|
||||
assert usage.total_tokens == 0
|
||||
|
||||
def test_with_values(self) -> None:
|
||||
"""Test with provided values."""
|
||||
usage = EmbeddingUsage(prompt_tokens=10, total_tokens=10)
|
||||
assert usage.prompt_tokens == 10
|
||||
assert usage.total_tokens == 10
|
||||
|
||||
|
||||
class TestEmbeddingData:
|
||||
"""Tests for EmbeddingData model."""
|
||||
|
||||
def test_creation(self) -> None:
|
||||
"""Test creating embedding data."""
|
||||
data = EmbeddingData(
|
||||
index=0,
|
||||
embedding=[0.1, 0.2, 0.3],
|
||||
)
|
||||
assert data.object == "embedding"
|
||||
assert data.index == 0
|
||||
assert data.embedding == [0.1, 0.2, 0.3]
|
||||
|
||||
def test_forbids_extra_fields(self) -> None:
|
||||
"""Test that extra fields are forbidden."""
|
||||
from pydantic import ValidationError
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
EmbeddingData(
|
||||
index=0,
|
||||
embedding=[0.1],
|
||||
extra_field="not allowed", # type: ignore[call-arg]
|
||||
)
|
||||
|
||||
|
||||
class TestModelInfo:
|
||||
"""Tests for ModelInfo model."""
|
||||
|
||||
def test_required_fields(self) -> None:
|
||||
"""Test required fields."""
|
||||
model = ModelInfo(id="test-model", backend="vllm")
|
||||
assert model.id == "test-model"
|
||||
assert model.backend == "vllm"
|
||||
|
||||
def test_optional_fields(self) -> None:
|
||||
"""Test optional fields with defaults."""
|
||||
model = ModelInfo(id="test-model", backend="vllm")
|
||||
assert model.loaded is True
|
||||
assert model.dimensions is None
|
||||
assert model.max_input_tokens is None
|
||||
|
||||
def test_with_all_fields(self) -> None:
|
||||
"""Test with all fields provided."""
|
||||
model = ModelInfo(
|
||||
id="test-model",
|
||||
backend="vllm",
|
||||
loaded=False,
|
||||
dimensions=768,
|
||||
max_input_tokens=8192,
|
||||
)
|
||||
assert model.loaded is False
|
||||
assert model.dimensions == 768
|
||||
assert model.max_input_tokens == 8192
|
||||
Loading…
Add table
Add a link
Reference in a new issue