Livrare LOT 1 - Didi
This commit is contained in:
commit
5380c3fc63
990 changed files with 133308 additions and 0 deletions
159
ai_platform/modules/llm-inference/.env.example
Normal file
159
ai_platform/modules/llm-inference/.env.example
Normal file
|
|
@ -0,0 +1,159 @@
|
|||
# LLM Inference 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: litellm, vllm, or llamacpp
|
||||
LLM_DEFAULT_BACKEND=litellm
|
||||
|
||||
# Enable optional backends (true/false)
|
||||
LLM_ENABLE_VLLM=false
|
||||
LLM_ENABLE_LLAMACPP=false
|
||||
|
||||
# =============================================================================
|
||||
# REQUIRED: API Keys (at least one for litellm backend)
|
||||
# =============================================================================
|
||||
|
||||
# OpenRouter API key (access 100+ models via single API)
|
||||
# Get one at: https://openrouter.ai/
|
||||
OPENROUTER_API_KEY=
|
||||
|
||||
# OpenAI API key
|
||||
# Get one at: https://platform.openai.com/
|
||||
OPENAI_API_KEY=
|
||||
|
||||
# Anthropic API key
|
||||
# Get one at: https://console.anthropic.com/
|
||||
ANTHROPIC_API_KEY=
|
||||
|
||||
# =============================================================================
|
||||
# REQUIRED for vLLM/llama.cpp profiles: Model Configuration
|
||||
# =============================================================================
|
||||
|
||||
# Directory containing model files (absolute path recommended)
|
||||
MODELS_DIR=/path/to/models
|
||||
|
||||
# vLLM model to load (Hugging Face model ID). Platform primary model (offer).
|
||||
# Start vLLM with `--served-model-name qwen3.5` so callers can use the alias.
|
||||
VLLM_MODEL=Qwen/Qwen3.5-35B-A3B
|
||||
|
||||
# llama.cpp model file (GGUF format, filename only - must be in MODELS_DIR)
|
||||
LLAMACPP_MODEL=qwen3.5-35b-a3b.Q4_K_M.gguf
|
||||
|
||||
# Default model alias used when a request omits `model`.
|
||||
LLM_DEFAULT_MODEL=qwen3.5
|
||||
|
||||
# Optional: map friendly alias -> real served model id (JSON). Leave unset if
|
||||
# the backend already serves the model under the alias name (e.g. via
|
||||
# vLLM --served-model-name qwen3.5).
|
||||
# LLM_MODEL_ALIASES={"qwen3.5":"Qwen/Qwen3.5-35B-A3B"}
|
||||
|
||||
# llama.cpp performance settings
|
||||
LLAMACPP_THREADS=4
|
||||
LLAMACPP_PARALLEL=1
|
||||
|
||||
# =============================================================================
|
||||
# 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
|
||||
# LLM_API_TOKENS=token1,token2,token3
|
||||
|
||||
# =============================================================================
|
||||
# REQUIRED: External URL (for OpenAPI spec and catalog integration)
|
||||
# =============================================================================
|
||||
|
||||
# External URL where this API is reachable (used in OpenAPI spec and /v1/info)
|
||||
LLM_EXTERNAL_URL=http://localhost:14011
|
||||
|
||||
# =============================================================================
|
||||
# OPTIONAL: API Server Settings
|
||||
# =============================================================================
|
||||
|
||||
# Server binding
|
||||
# LLM_HOST=0.0.0.0
|
||||
# LLM_PORT=14011
|
||||
|
||||
# Default model when not specified in request
|
||||
# LLM_DEFAULT_MODEL=gpt-3.5-turbo
|
||||
|
||||
# =============================================================================
|
||||
# OPTIONAL: Timeout and Retry Settings
|
||||
# =============================================================================
|
||||
|
||||
# Request timeout for LLM calls (seconds)
|
||||
# LLM_REQUEST_TIMEOUT=120.0
|
||||
|
||||
# Connection timeout (seconds)
|
||||
# LLM_CONNECT_TIMEOUT=10.0
|
||||
|
||||
# Maximum retry attempts for transient failures
|
||||
# LLM_MAX_RETRIES=3
|
||||
|
||||
# Minimum wait between retries (seconds)
|
||||
# LLM_RETRY_MIN_WAIT=1.0
|
||||
|
||||
# Maximum wait between retries (seconds)
|
||||
# LLM_RETRY_MAX_WAIT=60.0
|
||||
|
||||
# =============================================================================
|
||||
# OPTIONAL: Rate Limiting and Concurrency
|
||||
# =============================================================================
|
||||
|
||||
# Requests per second limit
|
||||
# LLM_RATE_LIMIT_RPS=10.0
|
||||
|
||||
# Maximum burst size for rate limiting
|
||||
# LLM_RATE_LIMIT_BURST=20
|
||||
|
||||
# Maximum concurrent completion requests
|
||||
# LLM_MAX_CONCURRENT_COMPLETIONS=10
|
||||
|
||||
# =============================================================================
|
||||
# OPTIONAL: Logging
|
||||
# =============================================================================
|
||||
|
||||
# Log level: DEBUG, INFO, WARNING, ERROR, CRITICAL
|
||||
# LLM_LOG_LEVEL=INFO
|
||||
|
||||
# Enable JSON logging format (true/false)
|
||||
# LLM_LOG_JSON=false
|
||||
|
||||
# =============================================================================
|
||||
# OPTIONAL: Nginx Proxy Timeouts (for Docker Compose deployment)
|
||||
# =============================================================================
|
||||
|
||||
# Nginx connection timeout
|
||||
# NGINX_CONNECT_TIMEOUT=60s
|
||||
|
||||
# Nginx send timeout
|
||||
# NGINX_SEND_TIMEOUT=120s
|
||||
|
||||
# Nginx read timeout (should be >= LLM_REQUEST_TIMEOUT)
|
||||
# NGINX_READ_TIMEOUT=600s
|
||||
|
||||
# =============================================================================
|
||||
# OPTIONAL: Backend-Specific URLs
|
||||
# =============================================================================
|
||||
|
||||
# vLLM server URL (internal Docker network)
|
||||
# LLM_VLLM_BASE_URL=http://didiAI-vllm-qwen3.5:14001
|
||||
|
||||
# llama.cpp server URL (single server, legacy)
|
||||
# LLM_LLAMACPP_BASE_URL=http://didiAI-llm-llamacpp:8080
|
||||
|
||||
# llama.cpp server URLs (comma-separated, load balanced round-robin with failover)
|
||||
# If set, overrides LLM_LLAMACPP_BASE_URL
|
||||
# LLM_LLAMACPP_BASE_URLS=http://10.11.10.43:14001,http://10.11.10.18:14001
|
||||
|
||||
# Health check interval for llama.cpp servers (seconds)
|
||||
# LLM_LLAMACPP_HEALTH_CHECK_INTERVAL=30
|
||||
|
||||
# vLLM API key (if authentication is enabled on vLLM server)
|
||||
# LLM_VLLM_API_KEY=
|
||||
611
ai_platform/modules/llm-inference/API.md
Normal file
611
ai_platform/modules/llm-inference/API.md
Normal file
|
|
@ -0,0 +1,611 @@
|
|||
# LLM Inference API Documentation
|
||||
|
||||
OpenAI-compatible REST API for unified LLM inference across multiple backends.
|
||||
|
||||
## Base URL
|
||||
|
||||
```
|
||||
{BASE_URL}
|
||||
```
|
||||
|
||||
Common configurations:
|
||||
- **Local development:** `http://localhost:14011`
|
||||
- **Docker (internal):** `http://didiAI-llm-api:14011`
|
||||
- **Docker (external):** `http://<host>:14011`
|
||||
- **VPN/Production:** Use your configured hostname or IP
|
||||
|
||||
## Authentication
|
||||
|
||||
### Bearer Token Authentication (Optional)
|
||||
|
||||
The API supports optional Bearer token authentication. When enabled, protected endpoints require a valid token.
|
||||
|
||||
**Enable by setting:**
|
||||
```bash
|
||||
LLM_API_TOKENS=token1,token2,token3 # comma-separated for multiple tokens
|
||||
```
|
||||
|
||||
**Request header:**
|
||||
```
|
||||
Authorization: Bearer <your-token>
|
||||
```
|
||||
|
||||
**Public endpoints (no auth required):**
|
||||
- `GET /health`
|
||||
- `GET /ready`
|
||||
|
||||
**Protected endpoints (require auth when enabled):**
|
||||
- `POST /v1/chat/completions`
|
||||
- `GET /v1/models`
|
||||
- `POST /v1/models/load`
|
||||
- `POST /v1/models/unload`
|
||||
- `GET /v1/backends`
|
||||
|
||||
**Error response (401):**
|
||||
```json
|
||||
{
|
||||
"detail": {
|
||||
"error": "Authentication required",
|
||||
"message": "Missing Authorization header"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Response includes `WWW-Authenticate: Bearer` header.
|
||||
|
||||
### Backend API Keys
|
||||
|
||||
Backend services require their own API keys:
|
||||
- **LiteLLM**: `OPENROUTER_API_KEY`, `OPENAI_API_KEY`, or `ANTHROPIC_API_KEY`
|
||||
- **vLLM**: Optional `LLM_VLLM_API_KEY` for vLLM server
|
||||
|
||||
## Rate Limiting
|
||||
|
||||
The API uses token bucket rate limiting:
|
||||
|
||||
| Header | Description |
|
||||
|--------|-------------|
|
||||
| `Retry-After` | Seconds to wait when rate limited (429 response) |
|
||||
| `X-Request-ID` | Unique request identifier (auto-generated or pass via header) |
|
||||
|
||||
Default limits:
|
||||
- **10 requests/second** with burst of 20
|
||||
- **10 concurrent completion requests**
|
||||
|
||||
Configure via environment variables:
|
||||
- `LLM_RATE_LIMIT_RPS`: Requests per second
|
||||
- `LLM_RATE_LIMIT_BURST`: Maximum burst size
|
||||
- `LLM_MAX_CONCURRENT_COMPLETIONS`: Concurrent request limit
|
||||
|
||||
---
|
||||
|
||||
## Endpoints
|
||||
|
||||
### Chat Completions
|
||||
|
||||
Create a chat completion with optional streaming.
|
||||
|
||||
```
|
||||
POST /v1/chat/completions
|
||||
```
|
||||
|
||||
#### Request Body
|
||||
|
||||
| Field | Type | Required | Default | Description |
|
||||
|-------|------|----------|---------|-------------|
|
||||
| `messages` | array | Yes | - | List of messages (1-1000) |
|
||||
| `model` | string | Yes | - | Model identifier |
|
||||
| `temperature` | float | No | 0.7 | Sampling temperature (0.0-2.0) |
|
||||
| `max_tokens` | integer | No | null | Maximum tokens to generate (1-1,000,000) |
|
||||
| `stream` | boolean | No | false | Enable streaming response |
|
||||
| `backend` | string | No | null | Backend override: `litellm`, `vllm`, `llamacpp` |
|
||||
| `top_p` | float | No | null | Top-p sampling (0.0-1.0) |
|
||||
| `frequency_penalty` | float | No | null | Frequency penalty (-2.0 to 2.0) |
|
||||
| `presence_penalty` | float | No | null | Presence penalty (-2.0 to 2.0) |
|
||||
| `stop` | string/array | No | null | Stop sequences |
|
||||
|
||||
#### Message Object
|
||||
|
||||
| Field | Type | Required | Description |
|
||||
|-------|------|----------|-------------|
|
||||
| `role` | string | Yes | One of: `system`, `user`, `assistant`, `function`, `tool` |
|
||||
| `content` | string | Yes* | Message content (*can be null for assistant role) |
|
||||
| `name` | string | No | Optional author name |
|
||||
|
||||
#### Response (Non-Streaming)
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "chatcmpl-abc123",
|
||||
"object": "chat.completion",
|
||||
"created": 1704067200,
|
||||
"model": "gpt-3.5-turbo",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "Hello! How can I help you today?"
|
||||
},
|
||||
"finish_reason": "stop"
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 15,
|
||||
"total_tokens": 25
|
||||
},
|
||||
"backend": "litellm"
|
||||
}
|
||||
```
|
||||
|
||||
#### Response (Streaming)
|
||||
|
||||
Server-Sent Events (SSE) format. Each event contains a JSON chunk:
|
||||
|
||||
```
|
||||
data: {"id":"chatcmpl-abc123","object":"chat.completion.chunk","created":1704067200,"model":"gpt-3.5-turbo","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}
|
||||
|
||||
data: {"id":"chatcmpl-abc123","object":"chat.completion.chunk","created":1704067200,"model":"gpt-3.5-turbo","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]}
|
||||
|
||||
data: {"id":"chatcmpl-abc123","object":"chat.completion.chunk","created":1704067200,"model":"gpt-3.5-turbo","choices":[{"index":0,"delta":{"content":"!"},"finish_reason":null}]}
|
||||
|
||||
data: {"id":"chatcmpl-abc123","object":"chat.completion.chunk","created":1704067200,"model":"gpt-3.5-turbo","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}
|
||||
|
||||
data: [DONE]
|
||||
```
|
||||
|
||||
#### Example Request
|
||||
|
||||
```bash
|
||||
# Non-streaming
|
||||
curl -X POST http://localhost:14011/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "What is the capital of France?"}
|
||||
],
|
||||
"model": "gpt-3.5-turbo",
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 100
|
||||
}'
|
||||
|
||||
# Streaming
|
||||
curl -X POST http://localhost:14011/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Accept: text/event-stream" \
|
||||
-d '{
|
||||
"messages": [{"role": "user", "content": "Tell me a short story"}],
|
||||
"model": "gpt-4",
|
||||
"stream": true
|
||||
}'
|
||||
|
||||
# With specific backend
|
||||
curl -X POST http://localhost:14011/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [{"role": "user", "content": "Hello!"}],
|
||||
"model": "meta-llama/Llama-2-7b-chat-hf",
|
||||
"backend": "vllm"
|
||||
}'
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### List Models
|
||||
|
||||
List available models across all or specific backends.
|
||||
|
||||
```
|
||||
GET /v1/models
|
||||
GET /v1/models?backend=litellm
|
||||
```
|
||||
|
||||
#### Query Parameters
|
||||
|
||||
| Parameter | Type | Required | Description |
|
||||
|-----------|------|----------|-------------|
|
||||
| `backend` | string | No | Filter by backend: `litellm`, `vllm`, `llamacpp` |
|
||||
|
||||
#### Response
|
||||
|
||||
```json
|
||||
{
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"id": "gpt-3.5-turbo",
|
||||
"backend": "litellm",
|
||||
"loaded": true,
|
||||
"context_length": 16384,
|
||||
"capabilities": ["chat"]
|
||||
},
|
||||
{
|
||||
"id": "gpt-4",
|
||||
"backend": "litellm",
|
||||
"loaded": true,
|
||||
"context_length": 128000,
|
||||
"capabilities": ["chat"]
|
||||
},
|
||||
{
|
||||
"id": "meta-llama/Llama-2-7b-chat-hf",
|
||||
"backend": "vllm",
|
||||
"loaded": true,
|
||||
"context_length": 4096,
|
||||
"capabilities": ["chat"]
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
#### Example Request
|
||||
|
||||
```bash
|
||||
# All models
|
||||
curl http://localhost:14011/v1/models
|
||||
|
||||
# Models from specific backend
|
||||
curl "http://localhost:14011/v1/models?backend=vllm"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Load Model
|
||||
|
||||
Load a model on a local backend (vLLM or llama.cpp only).
|
||||
|
||||
```
|
||||
POST /v1/models/load
|
||||
```
|
||||
|
||||
#### Request Body
|
||||
|
||||
| Field | Type | Required | Description |
|
||||
|-------|------|----------|-------------|
|
||||
| `model` | string | Yes | Model identifier to load |
|
||||
| `backend` | string | Yes | Target backend: `vllm` or `llamacpp` |
|
||||
|
||||
#### Response
|
||||
|
||||
```json
|
||||
{
|
||||
"success": true,
|
||||
"model": "meta-llama/Llama-2-7b-chat-hf",
|
||||
"backend": "vllm",
|
||||
"message": "Model loaded successfully"
|
||||
}
|
||||
```
|
||||
|
||||
#### Example Request
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:14011/v1/models/load \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "meta-llama/Llama-2-7b-chat-hf",
|
||||
"backend": "vllm"
|
||||
}'
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Unload Model
|
||||
|
||||
Unload a model from a local backend (vLLM or llama.cpp only).
|
||||
|
||||
```
|
||||
POST /v1/models/unload
|
||||
```
|
||||
|
||||
#### Request Body
|
||||
|
||||
| Field | Type | Required | Description |
|
||||
|-------|------|----------|-------------|
|
||||
| `model` | string | Yes | Model identifier to unload |
|
||||
| `backend` | string | Yes | Backend to unload from: `vllm` or `llamacpp` |
|
||||
|
||||
#### Response
|
||||
|
||||
```json
|
||||
{
|
||||
"success": true,
|
||||
"model": "meta-llama/Llama-2-7b-chat-hf",
|
||||
"backend": "vllm",
|
||||
"message": "Model unloaded successfully"
|
||||
}
|
||||
```
|
||||
|
||||
#### Example Request
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:14011/v1/models/unload \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "meta-llama/Llama-2-7b-chat-hf",
|
||||
"backend": "vllm"
|
||||
}'
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### List Backends
|
||||
|
||||
List available backends.
|
||||
|
||||
```
|
||||
GET /v1/backends
|
||||
```
|
||||
|
||||
#### Response
|
||||
|
||||
```json
|
||||
{
|
||||
"backends": ["litellm", "vllm", "llamacpp"]
|
||||
}
|
||||
```
|
||||
|
||||
#### Example Request
|
||||
|
||||
```bash
|
||||
curl http://localhost:14011/v1/backends
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Health Check
|
||||
|
||||
Check overall system health and per-backend status.
|
||||
|
||||
```
|
||||
GET /health
|
||||
```
|
||||
|
||||
#### Response
|
||||
|
||||
```json
|
||||
{
|
||||
"status": "healthy",
|
||||
"backends": [
|
||||
{
|
||||
"name": "litellm",
|
||||
"healthy": true,
|
||||
"message": null
|
||||
},
|
||||
{
|
||||
"name": "vllm",
|
||||
"healthy": true,
|
||||
"message": null
|
||||
},
|
||||
{
|
||||
"name": "llamacpp",
|
||||
"healthy": false,
|
||||
"message": "Connection refused"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
#### Status Values
|
||||
|
||||
| Status | Description |
|
||||
|--------|-------------|
|
||||
| `healthy` | All backends are healthy |
|
||||
| `degraded` | Some backends are healthy |
|
||||
| `unhealthy` | No backends are healthy |
|
||||
|
||||
#### Example Request
|
||||
|
||||
```bash
|
||||
curl http://localhost:14011/health
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Readiness Probe
|
||||
|
||||
Kubernetes-style readiness probe. Returns ready if the default backend is healthy.
|
||||
|
||||
```
|
||||
GET /ready
|
||||
```
|
||||
|
||||
#### Response
|
||||
|
||||
```json
|
||||
{
|
||||
"ready": true
|
||||
}
|
||||
```
|
||||
|
||||
#### HTTP Status Codes
|
||||
|
||||
| Code | Description |
|
||||
|------|-------------|
|
||||
| 200 | Service is ready |
|
||||
| 503 | Service is not ready |
|
||||
|
||||
#### Example Request
|
||||
|
||||
```bash
|
||||
curl http://localhost:14011/ready
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Error Responses
|
||||
|
||||
All errors follow a consistent format:
|
||||
|
||||
```json
|
||||
{
|
||||
"detail": "Error message description"
|
||||
}
|
||||
```
|
||||
|
||||
### HTTP Status Codes
|
||||
|
||||
| Code | Description |
|
||||
|------|-------------|
|
||||
| 400 | Bad Request - Invalid parameters or backend |
|
||||
| 401 | Unauthorized - Missing or invalid Bearer token (when auth enabled) |
|
||||
| 422 | Validation Error - Request body validation failed |
|
||||
| 429 | Too Many Requests - Rate limit exceeded |
|
||||
| 500 | Internal Server Error - Completion or backend failure |
|
||||
| 503 | Service Unavailable - Concurrency limit exceeded or service not ready |
|
||||
|
||||
### Error Examples
|
||||
|
||||
**Validation Error (422)**
|
||||
```json
|
||||
{
|
||||
"detail": [
|
||||
{
|
||||
"type": "value_error",
|
||||
"loc": ["body", "messages", 0, "content"],
|
||||
"msg": "Message at index 0 with role 'user' cannot have empty content",
|
||||
"input": ""
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
**Rate Limit (429)**
|
||||
```json
|
||||
{
|
||||
"detail": "Rate limit exceeded"
|
||||
}
|
||||
```
|
||||
Response includes `Retry-After` header with seconds to wait.
|
||||
|
||||
**Backend Error (400)**
|
||||
```json
|
||||
{
|
||||
"detail": "Backend 'vllm' is not available"
|
||||
}
|
||||
```
|
||||
|
||||
**Concurrency Limit (503)**
|
||||
```json
|
||||
{
|
||||
"detail": "Concurrency limit exceeded"
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Request Headers
|
||||
|
||||
| Header | Required | Description |
|
||||
|--------|----------|-------------|
|
||||
| `Content-Type` | Yes (POST) | Must be `application/json` |
|
||||
| `Authorization` | When auth enabled | Bearer token: `Bearer <your-token>` |
|
||||
| `Accept` | No | Use `text/event-stream` for streaming |
|
||||
| `X-Request-ID` | No | Custom request ID (auto-generated if not provided) |
|
||||
|
||||
---
|
||||
|
||||
## Supported Backends
|
||||
|
||||
### LiteLLM (Default)
|
||||
|
||||
Supports 100+ LLM providers through a unified interface.
|
||||
|
||||
**Popular models:**
|
||||
- `gpt-3.5-turbo`, `gpt-4`, `gpt-4-turbo` (OpenAI)
|
||||
- `claude-3-opus`, `claude-3-sonnet`, `claude-3-haiku` (Anthropic)
|
||||
- `openrouter/meta-llama/llama-3-70b` (OpenRouter)
|
||||
|
||||
### vLLM
|
||||
|
||||
High-throughput GPU inference for open-source models.
|
||||
|
||||
**Requirements:** NVIDIA GPU with 16GB+ VRAM
|
||||
|
||||
**Example models:**
|
||||
- `meta-llama/Llama-2-7b-chat-hf`
|
||||
- `mistralai/Mistral-7B-Instruct-v0.2`
|
||||
- `microsoft/phi-2`
|
||||
|
||||
### llama.cpp
|
||||
|
||||
CPU/Metal inference for GGUF format models.
|
||||
|
||||
**Requirements:** 16GB+ RAM, CPU with AVX2 support
|
||||
|
||||
**Example models:** Local GGUF files
|
||||
|
||||
---
|
||||
|
||||
## Python Client Example
|
||||
|
||||
```python
|
||||
import httpx
|
||||
|
||||
async def chat_completion():
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.post(
|
||||
"http://localhost:14011/v1/chat/completions",
|
||||
json={
|
||||
"messages": [{"role": "user", "content": "Hello!"}],
|
||||
"model": "gpt-3.5-turbo",
|
||||
},
|
||||
)
|
||||
return response.json()
|
||||
|
||||
async def streaming_completion():
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with client.stream(
|
||||
"POST",
|
||||
"http://localhost:14011/v1/chat/completions",
|
||||
json={
|
||||
"messages": [{"role": "user", "content": "Tell me a story"}],
|
||||
"model": "gpt-4",
|
||||
"stream": True,
|
||||
},
|
||||
) as response:
|
||||
async for line in response.aiter_lines():
|
||||
if line.startswith("data: "):
|
||||
data = line[6:]
|
||||
if data != "[DONE]":
|
||||
chunk = json.loads(data)
|
||||
content = chunk["choices"][0]["delta"].get("content", "")
|
||||
print(content, end="", flush=True)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## OpenAI SDK Compatibility
|
||||
|
||||
The API is compatible with the OpenAI Python SDK:
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
# Without auth (when LLM_API_TOKENS not set)
|
||||
client = OpenAI(
|
||||
base_url="http://localhost:14011/v1",
|
||||
api_key="not-needed", # Required by SDK but not used
|
||||
)
|
||||
|
||||
# With auth (when LLM_API_TOKENS is set)
|
||||
client = OpenAI(
|
||||
base_url="http://localhost:14011/v1",
|
||||
api_key="your-api-token", # Your LLM_API_TOKENS value
|
||||
)
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": "Hello!"}],
|
||||
)
|
||||
print(response.choices[0].message.content)
|
||||
|
||||
# Streaming
|
||||
for chunk in client.chat.completions.create(
|
||||
model="gpt-4",
|
||||
messages=[{"role": "user", "content": "Tell me a story"}],
|
||||
stream=True,
|
||||
):
|
||||
if chunk.choices[0].delta.content:
|
||||
print(chunk.choices[0].delta.content, end="")
|
||||
```
|
||||
163
ai_platform/modules/llm-inference/INDEX.md
Normal file
163
ai_platform/modules/llm-inference/INDEX.md
Normal file
|
|
@ -0,0 +1,163 @@
|
|||
# llm-inference
|
||||
|
||||
OpenAI-compatible LLM router for DIDI. Routes requests to local Qwen 3.5 397B (llama.cpp / vLLM backends) for free tier and proxies to OpenRouter cloud (Claude Sonnet 4.6, Gemini Flash 3, GPT-4o, etc.) for premium models. All endpoints are OpenAI-compatible (`/v1/chat/completions`), so callers can use the OpenAI SDK or plain `httpx` interchangeably.
|
||||
|
||||
- **Stack**: Python 3.10+, FastAPI, uvicorn, httpx, LiteLLM, Pydantic v2, sse-starlette
|
||||
- **URL**: `http://10.11.10.17:14011` (LLM router on GPU host) — referenced as `LLM_ROUTER_URL` in DIDI services
|
||||
- **Container**: `didiAI-llm-api` (image `didiai-llm-api`), runs on GPU host (typically `10.11.10.17`)
|
||||
- **Local model**: Qwen 3.5-35B-A3B (MoE, native multimodal text+vision) via vLLM (`didiAI-vllm-qwen3.5`, internal port 14001) or Qwen 3.5 397B-A17B variant via llama.cpp pool
|
||||
- **Entry point**: `llm-inference` console script -> `src/llm_inference/cli.py:main` -> uvicorn factory `llm_inference.api.app:create_app`
|
||||
|
||||
## Ce face
|
||||
|
||||
Single OpenAI-compatible endpoint (`/v1/chat/completions`) that selects a backend based on the request's `model` field and request-time `backend` override:
|
||||
|
||||
- **Auto backend resolution** (`LLMClient._resolve_backend_for_model` in `src/llm_inference/client.py`): probes each enabled local backend's `list_models()`. If the requested model is served locally, route to that backend; otherwise fall back to the configured default (litellm).
|
||||
- **Explicit backend override**: clients may pass `"backend": "litellm" | "vllm" | "llamacpp"` in the JSON body to force routing.
|
||||
- **Streaming + non-streaming**: same endpoint; `"stream": true` returns SSE chunks (`text/event-stream` with `data: [DONE]` terminator).
|
||||
- **Multimodal**: messages are pre-processed by `image_processing.process_messages` so image URLs/base64 attachments work the same across backends.
|
||||
- **Retries**: transient failures (`LLMTimeoutError`, `LLMRateLimitError`, `LLMConnectionError`) wrapped by `retry.retry_with_backoff` with exponential backoff (`max_retries`, `retry_min_wait`, `retry_max_wait`).
|
||||
- **Provider keys**: read from env (`OPENROUTER_API_KEY`, `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`); LiteLLM dispatches to OpenRouter / OpenAI / Anthropic / Gemini / Bedrock / 100+ providers using its own model-id parsing (`openrouter/...`, `claude-3-...`, `gpt-4o`, etc.).
|
||||
- **Optional Bearer auth**: when `LLM_API_TOKENS` is set, all `/v1/*` endpoints require `Authorization: Bearer <token>`; `/health` and `/ready` are always public.
|
||||
|
||||
## API endpoints
|
||||
|
||||
All paths are mounted by `src/llm_inference/api/app.py`:
|
||||
|
||||
| Method | Path | Purpose |
|
||||
|--------|------|---------|
|
||||
| POST | `/v1/chat/completions` | OpenAI-compat chat completion (streaming + non-streaming) |
|
||||
| GET | `/v1/models` | List models from all (or specific via `?backend=`) backends |
|
||||
| POST | `/v1/models/load` | Load a model on a local backend (vLLM / llama.cpp) |
|
||||
| POST | `/v1/models/unload` | Unload a model from a local backend |
|
||||
| GET | `/v1/backends` | List enabled backends |
|
||||
| GET | `/health` | Per-backend health (overall: `healthy` / `degraded` / `unhealthy`) |
|
||||
| GET | `/ready` | K8s-style readiness probe (200 if default backend healthy, 503 otherwise) |
|
||||
| GET | `/v1/info` *(catalog router)* | Service metadata for cross-module catalog integration |
|
||||
|
||||
Source files for the routers: `src/llm_inference/api/routes/{completions,models,health,info}.py`.
|
||||
|
||||
## Model routing logic
|
||||
|
||||
The backend dispatch lives in `LLMClient._resolve_backend_for_model` (auto) and `BackendRegistry.get` (explicit). Model identifier conventions:
|
||||
|
||||
- **Local vLLM** (`backend=vllm` or auto): the vLLM container exposes Qwen with `--served-model-name qwen3.5`. Requests with `model="qwen3.5"` (or whatever Hugging Face id is loaded) are forwarded to `LLM_VLLM_BASE_URL` (default `http://didiAI-vllm-qwen3.5:14001`) using OpenAI SDK.
|
||||
- **Local llama.cpp** (`backend=llamacpp` or auto): one or more llama.cpp servers (`LLM_LLAMACPP_BASE_URLS=http://10.11.10.43:14001,http://10.11.10.18:14001`) load-balanced round-robin with periodic health checks (`LLM_LLAMACPP_HEALTH_CHECK_INTERVAL`). See `src/llm_inference/backends/llamacpp_backend.py` (`LlamaCppBackend`, `LlamaCppServer` dataclass).
|
||||
- **Cloud / LiteLLM** (default): `model` is passed as-is to LiteLLM. Examples:
|
||||
- `openrouter/google/gemini-2.5-flash`, `openrouter/anthropic/claude-sonnet-4-6`, `openrouter/meta-llama/llama-3-70b` -> OpenRouter API (`OPENROUTER_API_KEY`)
|
||||
- `gpt-4o`, `gpt-4-turbo`, `gpt-3.5-turbo` -> OpenAI (`OPENAI_API_KEY`)
|
||||
- `claude-3-opus`, `claude-3-sonnet`, `claude-3-haiku` -> Anthropic (`ANTHROPIC_API_KEY`)
|
||||
- LiteLLM also handles Azure / Bedrock / Gemini natively — fallback chains can be implemented at the caller level by retrying with the next model id.
|
||||
|
||||
There is no Groq backend in the current code; cloud routing is consolidated through LiteLLM.
|
||||
|
||||
## How didi-brain + agent-v3 use it
|
||||
|
||||
This service is the single point of LLM accounting for the platform:
|
||||
|
||||
- **agent-v3** (`backend/services/orchestration-layer/agent-v3`): the `LLMClient.call()` helper inside `src/components/component-runner.ts` and the per-component executors (techniques, ai-tampered, claims, pipeline) call `${LLM_ROUTER_URL}/v1/chat/completions` for every analysis prompt. The `model` field is taken from the moderation config (Redis-synced from `didiFramework`) so admins can swap models without redeploying agent-v3.
|
||||
- **didi-brain** (`ai_platform/modules/didi_brain/brain_api`): the claim extractor and verification pipeline use the same router for both extraction prompts and rationale generation.
|
||||
- **All LLM traffic** from DIDI flows through this router — so any cost-tracking middleware added here gives a unified billing view.
|
||||
|
||||
## Structura fisiere
|
||||
|
||||
```
|
||||
src/llm_inference/
|
||||
__init__.py
|
||||
cli.py # `llm-inference` console script (uvicorn launcher + graceful shutdown)
|
||||
client.py # LLMClient (high-level Python API used as library)
|
||||
config.py # LLMSettings (Pydantic Settings, LLM_* env vars) + SettingsCache
|
||||
schemas.py # OpenAI-shaped CompletionResponse / CompletionChunk request/response models
|
||||
types.py # ChatMessage, BackendType enum, ModelInfo, Choice, Delta, Usage
|
||||
exceptions.py # CompletionError, LLMTimeoutError, LLMRateLimitError, LLMConnectionError, ModelLoadError
|
||||
retry.py # retry_with_backoff helper
|
||||
image_processing.py # process_messages — fetches/encodes image content for multimodal calls
|
||||
logging.py # configure_logging, get_logger (JSON + plain modes)
|
||||
utils.py # safe_close_stream, etc.
|
||||
backends/
|
||||
base.py # Abstract LLMBackend (complete / stream / list_models / health_check / load_model / unload_model)
|
||||
registry.py # BackendRegistry — instantiates enabled backends from settings
|
||||
litellm_backend.py # LiteLLMBackend (default; OpenRouter / OpenAI / Anthropic / 100+ via litellm)
|
||||
vllm_backend.py # vLLMBackend (OpenAI SDK -> internal vLLM server)
|
||||
llamacpp_backend.py # LlamaCppBackend (round-robin pool with health checks + failover)
|
||||
api/
|
||||
app.py # create_app() — FastAPI factory, lifespan, middleware wiring
|
||||
dependencies.py # init_concurrency_limiter, get_client, auth dependency
|
||||
middleware.py # RateLimitMiddleware (token bucket), RequestIdMiddleware
|
||||
routes/
|
||||
completions.py # POST /v1/chat/completions
|
||||
models.py # /v1/models, /v1/models/load, /v1/models/unload, /v1/backends
|
||||
health.py # /health, /ready
|
||||
info.py # /v1/info (catalog metadata)
|
||||
deploy/
|
||||
Dockerfile # Builds `didiai-llm-api` image
|
||||
docker-compose.yml # Profiles: api, vllm, llamacpp, full
|
||||
deploy.sh # Profile launcher
|
||||
nginx.conf # Optional reverse proxy
|
||||
tests/ # pytest (asyncio_mode=auto, markers: e2e, slow)
|
||||
API.md # OpenAPI-style endpoint docs
|
||||
README.md # Usage / install / deployment
|
||||
.env.example # Full env reference
|
||||
pyproject.toml # uv/hatchling project (llm-inference v0.1.0)
|
||||
```
|
||||
|
||||
## Configuration
|
||||
|
||||
All env vars use the `LLM_` prefix (Pydantic Settings, `extra="forbid"` so typos fail loudly). Required (no default):
|
||||
|
||||
- `LLM_DEFAULT_BACKEND` — `litellm` | `vllm` | `llamacpp`
|
||||
- `LLM_ENABLE_VLLM` — `true` / `false`
|
||||
- `LLM_ENABLE_LLAMACPP` — `true` / `false`
|
||||
- `LLM_EXTERNAL_URL` — public URL (used in OpenAPI `servers[]` and `/v1/info`)
|
||||
|
||||
Common optional:
|
||||
|
||||
- `OPENROUTER_API_KEY`, `OPENAI_API_KEY`, `ANTHROPIC_API_KEY` — provider keys for LiteLLM
|
||||
- `LLM_DEFAULT_MODEL` (default `gpt-3.5-turbo`)
|
||||
- `LLM_HOST` (default `0.0.0.0`), `LLM_PORT` (default `14011`)
|
||||
- `LLM_VLLM_BASE_URL` (default `http://localhost:14001`), `LLM_VLLM_API_KEY`
|
||||
- `LLM_LLAMACPP_BASE_URL` (single-server legacy) **or** `LLM_LLAMACPP_BASE_URLS` (comma-separated pool, overrides single)
|
||||
- `LLM_LLAMACPP_HEALTH_CHECK_INTERVAL` (default `30s`)
|
||||
- `LLM_REQUEST_TIMEOUT` (default `120s`), `LLM_CONNECT_TIMEOUT` (default `10s`)
|
||||
- `LLM_MAX_RETRIES` (default `3`), `LLM_RETRY_MIN_WAIT`, `LLM_RETRY_MAX_WAIT`
|
||||
- `LLM_RATE_LIMIT_RPS` (default `10`), `LLM_RATE_LIMIT_BURST` (default `20`), `LLM_MAX_CONCURRENT_COMPLETIONS` (default `10`)
|
||||
- `LLM_API_TOKENS` — comma-separated; if set, `/v1/*` requires `Authorization: Bearer <token>`
|
||||
- `LLM_LOG_LEVEL`, `LLM_LOG_JSON`
|
||||
|
||||
Full reference: `.env.example`. Settings class: `src/llm_inference/config.py:LLMSettings`.
|
||||
|
||||
## Deployment
|
||||
|
||||
GPU host `10.11.10.17` (typically). The vLLM container needs an NVIDIA H200 NVL (~143 GB VRAM); Qwen 3.5-35B-A3B occupies ~57 GB BF16 with `--gpu-memory-utilization 0.65` and `--max-model-len 32000`, leaving room on GPU 0 for Whisper.
|
||||
|
||||
Compose profiles (in `deploy/docker-compose.yml`):
|
||||
|
||||
- `api` — only the FastAPI router (uses external/cloud providers)
|
||||
- `vllm` — router + vLLM Qwen 3.5 (GPU)
|
||||
- `llamacpp` — router + a local llama.cpp server (CPU)
|
||||
- `full` — everything
|
||||
|
||||
Networking: shares the external Docker network `didi-network` with sibling AI modules (whisper, embeddings, etc.). Internal address used by the router for vLLM is `http://didiAI-vllm-qwen3.5:14001`.
|
||||
|
||||
```bash
|
||||
cd deploy/
|
||||
cp ../.env.example .env # fill keys + URLs
|
||||
./deploy.sh --profile api -d
|
||||
./deploy.sh --profile vllm -d
|
||||
./deploy.sh --profile api --logs
|
||||
./deploy.sh --profile api --down
|
||||
```
|
||||
|
||||
Healthcheck for the API container: `python -c "import urllib.request; urllib.request.urlopen('http://localhost:14011/health')"` every 30s.
|
||||
|
||||
## Cost tracking
|
||||
|
||||
The current code base does not yet ship a per-request cost-accounting middleware. The wiring point exists (`api/middleware.py` and the `RequestIdMiddleware` propagates a `X-Request-ID` to every call) and prompt/completion token counts are returned in `Usage` (`types.py`) on every non-streaming completion and in the final stream chunk. To plug into the AI dashboard, intercept at `routes/completions.py` (or a new middleware) and persist `(request_id, backend, model, prompt_tokens, completion_tokens, latency_ms)` — the dashboard's `cost.html` consumes that schema.
|
||||
|
||||
## Related
|
||||
|
||||
- **agent-v3 client**: `backend/services/orchestration-layer/agent-v3/src/components/component-runner.ts` — `LLMClient.call()` posts to `${LLM_ROUTER_URL}/v1/chat/completions`. Model id comes from Redis moderation config synced by `didiFramework`.
|
||||
- **didi-brain**: `ai_platform/modules/didi_brain/brain_api` — extractor + verification calls hit the same router.
|
||||
- **AI dashboard**: `cost.html` (when accounting middleware is added) consumes per-request token usage.
|
||||
- **Sibling AI modules** (sharing `didi-network` network): whisper / transcription, embeddings, vision, etc., living under `ai_platform/modules/`.
|
||||
- **OpenAI SDK compatibility**: clients can simply do `OpenAI(base_url="http://10.11.10.17:14011/v1", api_key=<token-or-not-needed>)`.
|
||||
235
ai_platform/modules/llm-inference/README.md
Normal file
235
ai_platform/modules/llm-inference/README.md
Normal file
|
|
@ -0,0 +1,235 @@
|
|||
# LLM Inference
|
||||
|
||||
Unified LLM inference module with multiple backends: LiteLLM (OpenRouter, OpenAI, Anthropic), vLLM, and llama.cpp.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
**Required:**
|
||||
- All global prerequisites (see main [README.md](../../README.md))
|
||||
- At least one LLM API key (OpenRouter, OpenAI, or Anthropic)
|
||||
|
||||
**For vLLM backend (GPU inference):**
|
||||
- NVIDIA GPU with 16GB+ VRAM (24GB+ recommended for larger models)
|
||||
- NVIDIA Driver 535+
|
||||
- NVIDIA Container Toolkit
|
||||
|
||||
**For llama.cpp backend (CPU inference):**
|
||||
- 16GB+ RAM (depends on model size)
|
||||
- CPU with AVX2 support (most modern CPUs)
|
||||
|
||||
## Features
|
||||
|
||||
- **Multiple Backends**: LiteLLM (100+ providers), vLLM (GPU inference), llama.cpp (CPU/Metal)
|
||||
- **Unified Interface**: Single API for all backends, switch seamlessly
|
||||
- **Streaming Support**: Server-Sent Events (SSE) for real-time responses
|
||||
- **Model Management**: List, load, and unload models (local backends)
|
||||
- **OpenAI Compatible**: Drop-in replacement for OpenAI API clients
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
cd modules/llm-inference
|
||||
|
||||
# Install core dependencies
|
||||
uv sync
|
||||
|
||||
# Install with optional backends
|
||||
uv sync --extra vllm # For vLLM support
|
||||
uv sync --extra llamacpp # For llama.cpp support
|
||||
uv sync --extra local # For all local backends
|
||||
uv sync --extra dev # For development
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
### As a Python Library
|
||||
|
||||
```python
|
||||
from llm_inference import LLMClient, BackendType
|
||||
|
||||
# Initialize client
|
||||
client = LLMClient()
|
||||
|
||||
# Simple completion
|
||||
response = await client.complete(
|
||||
messages=[{"role": "user", "content": "Hello!"}],
|
||||
model="gpt-3.5-turbo",
|
||||
)
|
||||
print(response.choices[0].message.content)
|
||||
|
||||
# Streaming
|
||||
async for chunk in client.stream(
|
||||
messages=[{"role": "user", "content": "Tell me a story"}],
|
||||
model="gpt-4",
|
||||
):
|
||||
if chunk.choices[0].delta.content:
|
||||
print(chunk.choices[0].delta.content, end="")
|
||||
|
||||
# Use specific backend
|
||||
response = await client.complete(
|
||||
messages=[{"role": "user", "content": "Hello!"}],
|
||||
model="meta-llama/Llama-2-7b-chat-hf",
|
||||
backend=BackendType.VLLM,
|
||||
)
|
||||
```
|
||||
|
||||
### As an API Server
|
||||
|
||||
```bash
|
||||
cd deploy/
|
||||
|
||||
# Copy and configure environment
|
||||
cp ../.env.example .env
|
||||
# Edit .env with your API keys and settings
|
||||
|
||||
# Start the server
|
||||
./deploy.sh --profile api -d
|
||||
```
|
||||
|
||||
### API Endpoints
|
||||
|
||||
| Endpoint | Method | Description |
|
||||
|----------|--------|-------------|
|
||||
| `/v1/chat/completions` | POST | Chat completion (supports streaming) |
|
||||
| `/v1/models` | GET | List available models |
|
||||
| `/v1/models/load` | POST | Load a model (local backends) |
|
||||
| `/v1/models/unload` | POST | Unload a model (local backends) |
|
||||
| `/v1/backends` | GET | List available backends |
|
||||
| `/health` | GET | Health check |
|
||||
| `/ready` | GET | Readiness probe |
|
||||
|
||||
### Example API Request
|
||||
|
||||
```bash
|
||||
# Non-streaming
|
||||
curl -X POST http://localhost:14011/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [{"role": "user", "content": "Hello!"}],
|
||||
"model": "gpt-3.5-turbo"
|
||||
}'
|
||||
|
||||
# Streaming
|
||||
curl -X POST http://localhost:14011/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [{"role": "user", "content": "Tell me a story"}],
|
||||
"model": "gpt-4",
|
||||
"stream": true
|
||||
}'
|
||||
```
|
||||
|
||||
## Configuration
|
||||
|
||||
Configure via environment variables (prefix: `LLM_`):
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `LLM_DEFAULT_BACKEND` | `litellm` | Default backend (litellm, vllm, llamacpp) |
|
||||
| `LLM_DEFAULT_MODEL` | `gpt-3.5-turbo` | Default model |
|
||||
| `LLM_PORT` | `14011` | API server port |
|
||||
| `LLM_HOST` | `0.0.0.0` | API server host |
|
||||
| `LLM_ENABLE_VLLM` | `false` | Enable vLLM backend |
|
||||
| `LLM_ENABLE_LLAMACPP` | `false` | Enable llama.cpp backend |
|
||||
| `LLM_VLLM_BASE_URL` | `http://localhost:14001` | vLLM server URL |
|
||||
| `LLM_LLAMACPP_BASE_URL` | `http://localhost:8080` | llama.cpp server URL |
|
||||
| `OPENROUTER_API_KEY` | - | OpenRouter API key |
|
||||
| `OPENAI_API_KEY` | - | OpenAI API key |
|
||||
| `ANTHROPIC_API_KEY` | - | Anthropic API key |
|
||||
|
||||
## Rate Limiting
|
||||
|
||||
The API server includes built-in rate limiting using a token bucket algorithm. Configure via environment variables:
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `LLM_RATE_LIMIT_RPS` | `10.0` | Requests per second |
|
||||
| `LLM_RATE_LIMIT_BURST` | `20` | Maximum burst size |
|
||||
| `LLM_MAX_CONCURRENT_COMPLETIONS` | `10` | Maximum concurrent completion requests |
|
||||
|
||||
**Important:** Rate limiting is **per-process**. In multi-replica deployments (e.g., Kubernetes), each replica has its own independent limit. For distributed rate limiting, use an external solution like Redis, an API gateway (Kong, nginx), or cloud provider rate limiting.
|
||||
|
||||
## Deployment
|
||||
|
||||
```bash
|
||||
cd deploy/
|
||||
|
||||
# Copy and configure environment (REQUIRED)
|
||||
cp ../.env.example .env
|
||||
# Edit .env with your settings
|
||||
|
||||
# API only (uses external LLM providers)
|
||||
./deploy.sh --profile api -d
|
||||
|
||||
# With vLLM (requires NVIDIA GPU)
|
||||
./deploy.sh --profile vllm -d
|
||||
|
||||
# With llama.cpp (CPU inference)
|
||||
./deploy.sh --profile llamacpp -d
|
||||
|
||||
# Full stack
|
||||
./deploy.sh --profile full -d
|
||||
|
||||
# View logs
|
||||
./deploy.sh --profile api --logs
|
||||
|
||||
# Stop services
|
||||
./deploy.sh --profile api --down
|
||||
```
|
||||
|
||||
### Port Allocation
|
||||
|
||||
| Port | Service |
|
||||
|------|---------|
|
||||
| 14011 | LLM Inference API |
|
||||
| 14001 | vLLM Qwen3.5-35B-A3B |
|
||||
|
||||
## Development
|
||||
|
||||
```bash
|
||||
# Install dev dependencies
|
||||
uv sync --extra dev
|
||||
|
||||
# Run tests
|
||||
uv run pytest
|
||||
|
||||
# Run tests with coverage
|
||||
uv run pytest --cov=src/llm_inference --cov-report=term-missing
|
||||
|
||||
# Lint and format
|
||||
uv run ruff check .
|
||||
uv run ruff format .
|
||||
|
||||
# Type check
|
||||
uv run mypy src/
|
||||
```
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
src/llm_inference/
|
||||
├── __init__.py # Package exports
|
||||
├── config.py # Pydantic Settings
|
||||
├── types.py # Core types and enums
|
||||
├── schemas.py # API request/response schemas
|
||||
├── exceptions.py # Custom exceptions
|
||||
├── client.py # High-level LLMClient
|
||||
├── cli.py # CLI entry point
|
||||
├── backends/
|
||||
│ ├── base.py # Abstract LLMBackend
|
||||
│ ├── registry.py # BackendRegistry
|
||||
│ ├── litellm_backend.py
|
||||
│ ├── vllm_backend.py
|
||||
│ └── llamacpp_backend.py
|
||||
└── api/
|
||||
├── app.py # FastAPI app factory
|
||||
├── dependencies.py
|
||||
└── routes/
|
||||
├── completions.py
|
||||
├── models.py
|
||||
└── health.py
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
43
ai_platform/modules/llm-inference/deploy/Dockerfile
Normal file
43
ai_platform/modules/llm-inference/deploy/Dockerfile
Normal file
|
|
@ -0,0 +1,43 @@
|
|||
# LLM Inference 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 uv.lock README.md ./
|
||||
COPY src/ ./src/
|
||||
|
||||
# Install dependencies
|
||||
RUN uv sync --frozen --no-dev
|
||||
|
||||
# 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 LLM_HOST=0.0.0.0
|
||||
ENV LLM_PORT=14011
|
||||
|
||||
# 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:14011/health')" || exit 1
|
||||
|
||||
# Expose port
|
||||
EXPOSE 14011
|
||||
|
||||
# Run the server
|
||||
CMD ["python", "-m", "llm_inference.cli"]
|
||||
154
ai_platform/modules/llm-inference/deploy/deploy.sh
Normal file
154
ai_platform/modules/llm-inference/deploy/deploy.sh
Normal file
|
|
@ -0,0 +1,154 @@
|
|||
#!/usr/bin/env bash
|
||||
#
|
||||
# Docker Compose Startup Script for LLM Inference
|
||||
#
|
||||
# Usage: ./deploy/docker-start.sh [OPTIONS]
|
||||
#
|
||||
# Options:
|
||||
# --profile <api|vllm|llamacpp|full> 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,16p' "$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, full"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Check required variables based on profile
|
||||
check_required_var "LLM_DEFAULT_BACKEND"
|
||||
check_required_var "LLM_ENABLE_VLLM"
|
||||
check_required_var "LLM_ENABLE_LLAMACPP"
|
||||
check_required_var "LLM_EXTERNAL_URL"
|
||||
|
||||
# API keys - at least one MUST be set for litellm backend (fail-fast)
|
||||
if [[ "$LLM_DEFAULT_BACKEND" == "litellm" ]]; then
|
||||
if [[ -z "${OPENROUTER_API_KEY:-}" && -z "${OPENAI_API_KEY:-}" && -z "${ANTHROPIC_API_KEY:-}" ]]; then
|
||||
echo "ERROR: litellm backend requires at least one API key"
|
||||
echo "Set one of: OPENROUTER_API_KEY, OPENAI_API_KEY, or ANTHROPIC_API_KEY"
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
|
||||
# Check model-specific variables for vllm/llamacpp profiles
|
||||
case $PROFILE in
|
||||
vllm|full)
|
||||
check_required_var "MODELS_DIR"
|
||||
check_required_var "VLLM_MODEL"
|
||||
;;
|
||||
llamacpp|full)
|
||||
check_required_var "MODELS_DIR"
|
||||
check_required_var "LLAMACPP_MODEL"
|
||||
check_required_var "LLAMACPP_THREADS"
|
||||
check_required_var "LLAMACPP_PARALLEL"
|
||||
;;
|
||||
esac
|
||||
|
||||
cd "$SCRIPT_DIR"
|
||||
|
||||
case $ACTION in
|
||||
up)
|
||||
echo "Starting LLM Inference with profile: $PROFILE"
|
||||
echo " Default backend: $LLM_DEFAULT_BACKEND"
|
||||
echo " vLLM enabled: $LLM_ENABLE_VLLM"
|
||||
echo " llama.cpp enabled: $LLM_ENABLE_LLAMACPP"
|
||||
if [[ "$PROFILE" == "vllm" || "$PROFILE" == "full" ]]; then
|
||||
echo " vLLM model: $VLLM_MODEL"
|
||||
fi
|
||||
if [[ "$PROFILE" == "llamacpp" || "$PROFILE" == "full" ]]; then
|
||||
echo " llama.cpp model: $LLAMACPP_MODEL"
|
||||
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 LLM Inference 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
|
||||
142
ai_platform/modules/llm-inference/deploy/docker-compose.yml
Normal file
142
ai_platform/modules/llm-inference/deploy/docker-compose.yml
Normal file
|
|
@ -0,0 +1,142 @@
|
|||
# LLM Inference Module - Docker Compose Configuration
|
||||
#
|
||||
# Port Allocation:
|
||||
# 14001 - vLLM Qwen3.5-35B-A3B (text + vision MoE, native multimodal)
|
||||
# 14011 - LLM Inference API (unified router)
|
||||
#
|
||||
# Profiles:
|
||||
# api - API server only (uses external LLM services)
|
||||
# vllm - API + vLLM servers (GPU required)
|
||||
#
|
||||
# Model:
|
||||
# - Qwen3.5-35B-A3B (~57GB weights BF16, gpu-util 0.45)
|
||||
# - Native vision support (replaces separate Qwen3-VL)
|
||||
#
|
||||
# Hardware: 1x NVIDIA H200 NVL (~143GB VRAM)
|
||||
# - GPU 0: Qwen3.5-35B-A3B + Whisper (~2GB)
|
||||
#
|
||||
# Naming Convention: didiAI-{module}-{service}
|
||||
#
|
||||
# Network:
|
||||
# Uses deploy_default network (shared with other modules)
|
||||
|
||||
networks:
|
||||
deploy_default:
|
||||
external: true
|
||||
|
||||
services:
|
||||
# ==========================================================================
|
||||
# LLM Inference API Server
|
||||
# ==========================================================================
|
||||
llm-api:
|
||||
container_name: didiAI-llm-api
|
||||
image: didiai-llm-api
|
||||
build:
|
||||
context: ..
|
||||
dockerfile: deploy/Dockerfile
|
||||
ports:
|
||||
- "14011:14011"
|
||||
networks:
|
||||
- deploy_default
|
||||
environment:
|
||||
- LLM_PORT=14011
|
||||
- LLM_EXTERNAL_URL=${LLM_EXTERNAL_URL}
|
||||
- LLM_DEFAULT_BACKEND=${LLM_DEFAULT_BACKEND}
|
||||
- LLM_ENABLE_VLLM=${LLM_ENABLE_VLLM}
|
||||
- LLM_ENABLE_LLAMACPP=${LLM_ENABLE_LLAMACPP}
|
||||
- LLM_VLLM_BASE_URL=http://didiAI-vllm-qwen3.5:14001
|
||||
- LLM_LLAMACPP_BASE_URLS=${LLM_LLAMACPP_BASE_URLS:-}
|
||||
- LLM_API_TOKENS=${LLM_API_TOKENS:-}
|
||||
- LLM_DASHBOARD_URL=${LLM_DASHBOARD_URL:-http://didiAI-dashboard:51300}
|
||||
- OPENROUTER_API_KEY=${OPENROUTER_API_KEY:-}
|
||||
- OPENAI_API_KEY=${OPENAI_API_KEY:-}
|
||||
- ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY:-}
|
||||
healthcheck:
|
||||
test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:14011/health')"]
|
||||
interval: 30s
|
||||
timeout: 10s
|
||||
retries: 3
|
||||
start_period: 30s
|
||||
restart: unless-stopped
|
||||
profiles:
|
||||
- api
|
||||
- vllm
|
||||
|
||||
# ==========================================================================
|
||||
# vLLM Server - Qwen3.5-35B-A3B (Unified Text + Vision MoE Model)
|
||||
# ==========================================================================
|
||||
# Native multimodal - replaces separate Qwen3 text + Qwen3-VL vision
|
||||
vllm-qwen3.5:
|
||||
container_name: didiAI-vllm-qwen3.5
|
||||
image: vllm/vllm-openai:qwen3_5
|
||||
ports:
|
||||
- "14001:14001"
|
||||
networks:
|
||||
- deploy_default
|
||||
volumes:
|
||||
- ${HF_CACHE_DIR}:/root/.cache/huggingface
|
||||
environment:
|
||||
- HF_HOME=/root/.cache/huggingface
|
||||
- HUGGING_FACE_HUB_TOKEN=${HF_TOKEN:-}
|
||||
- CUDA_VISIBLE_DEVICES=0
|
||||
- VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
|
||||
command: >
|
||||
--model Qwen/Qwen3.5-35B-A3B
|
||||
--host 0.0.0.0
|
||||
--port 14001
|
||||
--served-model-name qwen3.5
|
||||
--tensor-parallel-size 1
|
||||
--max-model-len 32000
|
||||
--gpu-memory-utilization 0.65
|
||||
--trust-remote-code
|
||||
--enable-prefix-caching
|
||||
--disable-log-requests
|
||||
--enable-auto-tool-choice
|
||||
--tool-call-parser hermes
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
device_ids: ['0']
|
||||
capabilities: [gpu]
|
||||
healthcheck:
|
||||
test: ["CMD", "curl", "-f", "http://localhost:14001/health"]
|
||||
interval: 30s
|
||||
timeout: 10s
|
||||
retries: 10
|
||||
start_period: 600s
|
||||
restart: unless-stopped
|
||||
profiles:
|
||||
- vllm
|
||||
|
||||
# ==========================================================================
|
||||
# llama.cpp Server (CPU/Metal Inference) - Optional
|
||||
# ==========================================================================
|
||||
llamacpp:
|
||||
image: ghcr.io/ggml-org/llama.cpp:server-b4769
|
||||
expose:
|
||||
- "14011"
|
||||
volumes:
|
||||
- ${MODELS_DIR:-/cai2_ds_storage/models}:/models:ro
|
||||
command: >
|
||||
--model /models/${LLAMACPP_MODEL:-model.gguf}
|
||||
--host 0.0.0.0
|
||||
--port 14011
|
||||
--ctx-size 4096
|
||||
--threads ${LLAMACPP_THREADS:-4}
|
||||
--parallel ${LLAMACPP_PARALLEL:-1}
|
||||
healthcheck:
|
||||
test: ["CMD", "curl", "-f", "http://localhost:14011/health"]
|
||||
interval: 30s
|
||||
timeout: 10s
|
||||
retries: 3
|
||||
start_period: 60s
|
||||
restart: unless-stopped
|
||||
profiles:
|
||||
- llamacpp
|
||||
- full
|
||||
|
||||
volumes:
|
||||
vllm-cache:
|
||||
name: llm-inference-vllm-cache
|
||||
49
ai_platform/modules/llm-inference/deploy/nginx.conf
Normal file
49
ai_platform/modules/llm-inference/deploy/nginx.conf
Normal file
|
|
@ -0,0 +1,49 @@
|
|||
upstream llm_api {
|
||||
server llm-api:8100;
|
||||
keepalive 32;
|
||||
}
|
||||
|
||||
server {
|
||||
listen 80;
|
||||
server_name _;
|
||||
|
||||
# Timeouts for slow LLM responses
|
||||
proxy_connect_timeout 60s;
|
||||
proxy_send_timeout 120s;
|
||||
proxy_read_timeout 300s;
|
||||
|
||||
# Health checks (no logging)
|
||||
location /health {
|
||||
access_log off;
|
||||
proxy_pass http://llm_api/health;
|
||||
proxy_http_version 1.1;
|
||||
proxy_set_header Host $host;
|
||||
}
|
||||
|
||||
location /ready {
|
||||
access_log off;
|
||||
proxy_pass http://llm_api/ready;
|
||||
proxy_http_version 1.1;
|
||||
proxy_set_header Host $host;
|
||||
}
|
||||
|
||||
# API endpoints
|
||||
location / {
|
||||
proxy_pass http://llm_api;
|
||||
proxy_http_version 1.1;
|
||||
|
||||
# Forward client info
|
||||
proxy_set_header Host $host;
|
||||
proxy_set_header X-Real-IP $remote_addr;
|
||||
proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
|
||||
|
||||
# SSE streaming support
|
||||
proxy_set_header Connection '';
|
||||
proxy_buffering off;
|
||||
proxy_cache off;
|
||||
chunked_transfer_encoding off;
|
||||
|
||||
# Large request bodies for long conversations
|
||||
client_max_body_size 10M;
|
||||
}
|
||||
}
|
||||
82
ai_platform/modules/llm-inference/pyproject.toml
Normal file
82
ai_platform/modules/llm-inference/pyproject.toml
Normal file
|
|
@ -0,0 +1,82 @@
|
|||
[project]
|
||||
name = "llm-inference"
|
||||
version = "0.1.0"
|
||||
description = "Unified LLM inference with multiple backends (LiteLLM, vLLM, llama.cpp)"
|
||||
requires-python = ">=3.10"
|
||||
readme = "README.md"
|
||||
|
||||
dependencies = [
|
||||
# Core dependencies (always installed)
|
||||
"litellm>=1.50.0,<2.0",
|
||||
"fastapi>=0.115.0,<0.116",
|
||||
"uvicorn[standard]>=0.32.0",
|
||||
"pydantic>=2.0,<3.0",
|
||||
"pydantic-settings>=2.0,<3.0",
|
||||
"httpx>=0.27.0,<1.0",
|
||||
"sse-starlette>=2.0,<3.0",
|
||||
"prometheus-client>=0.20.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 (GPU inference via OpenAI-compatible API)
|
||||
vllm = [
|
||||
"openai>=1.50.0,<2.0",
|
||||
]
|
||||
|
||||
# llama.cpp backend (CPU/Metal inference via OpenAI-compatible API)
|
||||
llamacpp = [
|
||||
"openai>=1.50.0,<2.0",
|
||||
]
|
||||
|
||||
# All local inference backends
|
||||
local = [
|
||||
"llm-inference[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]
|
||||
llm-inference = "llm_inference.cli:main"
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["src/llm_inference"]
|
||||
|
||||
[tool.ruff]
|
||||
extend = "../../ruff.toml"
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
asyncio_mode = "auto"
|
||||
testpaths = ["tests"]
|
||||
addopts = "-v --tb=short"
|
||||
markers = [
|
||||
"e2e: End-to-end tests against real running server",
|
||||
"slow: Tests that make real LLM API calls (cost money, slower)",
|
||||
]
|
||||
|
||||
[tool.mypy]
|
||||
python_version = "3.10"
|
||||
strict = true
|
||||
warn_return_any = true
|
||||
warn_unused_ignores = true
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"pytest>=9.0.2",
|
||||
"pytest-asyncio>=1.3.0",
|
||||
"ruff>=0.14.11",
|
||||
]
|
||||
|
|
@ -0,0 +1,16 @@
|
|||
"""LLM Inference - Unified LLM inference with multiple backends."""
|
||||
|
||||
from llm_inference.client import LLMClient
|
||||
from llm_inference.config import LLMSettings, get_settings
|
||||
from llm_inference.types import BackendType, ChatMessage
|
||||
|
||||
__version__ = "0.1.0"
|
||||
|
||||
__all__ = [
|
||||
"BackendType",
|
||||
"ChatMessage",
|
||||
"LLMClient",
|
||||
"LLMSettings",
|
||||
"__version__",
|
||||
"get_settings",
|
||||
]
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
"""FastAPI application for LLM inference."""
|
||||
|
||||
from llm_inference.api.app import create_app
|
||||
|
||||
__all__ = ["create_app"]
|
||||
165
ai_platform/modules/llm-inference/src/llm_inference/api/app.py
Normal file
165
ai_platform/modules/llm-inference/src/llm_inference/api/app.py
Normal file
|
|
@ -0,0 +1,165 @@
|
|||
"""FastAPI application factory."""
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
from llm_inference.api.dependencies import init_concurrency_limiter
|
||||
from llm_inference.api.middleware import RateLimitMiddleware, RequestIdMiddleware
|
||||
from llm_inference.api.routes import completions, health, info, models
|
||||
from llm_inference.backends.llamacpp_backend import LlamaCppBackend
|
||||
from llm_inference.client import LLMClient
|
||||
from llm_inference.config import SettingsCache
|
||||
from llm_inference.image_processing import close_http_client
|
||||
from llm_inference.logging import configure_logging, get_logger
|
||||
from llm_inference.runtime_config import RuntimeConfigClient
|
||||
from llm_inference.types import BackendType
|
||||
|
||||
logger = get_logger("app")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
|
||||
"""Application lifespan manager.
|
||||
|
||||
Initializes the LLM client, settings, logging, and concurrency limiter
|
||||
on startup, and performs cleanup on shutdown.
|
||||
|
||||
Args:
|
||||
app: FastAPI application instance.
|
||||
|
||||
Yields:
|
||||
None: Control to the application.
|
||||
"""
|
||||
settings = SettingsCache.get()
|
||||
configure_logging(settings.log_level, settings.log_json)
|
||||
|
||||
logger.info("Starting LLM Inference 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_completions)
|
||||
|
||||
app.state.settings = settings
|
||||
app.state.client = LLMClient(settings)
|
||||
|
||||
# Start the runtime config client created in create_app()
|
||||
if hasattr(app.state, "runtime_config"):
|
||||
await app.state.runtime_config.start()
|
||||
|
||||
if settings.enable_llamacpp:
|
||||
try:
|
||||
backend = app.state.client.registry.get(BackendType.LLAMACPP)
|
||||
if isinstance(backend, LlamaCppBackend):
|
||||
backend.start_health_checks()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
logger.info("LLM Inference API started successfully")
|
||||
|
||||
yield
|
||||
|
||||
logger.info("Shutting down LLM Inference API")
|
||||
if hasattr(app.state, "runtime_config"):
|
||||
await app.state.runtime_config.stop()
|
||||
await close_http_client()
|
||||
|
||||
|
||||
def create_app() -> FastAPI:
|
||||
"""Create and configure the FastAPI application.
|
||||
|
||||
Returns:
|
||||
FastAPI: Configured FastAPI application.
|
||||
"""
|
||||
settings = SettingsCache.get()
|
||||
|
||||
app = FastAPI(
|
||||
title="LLM Inference API",
|
||||
description=(
|
||||
"Unified LLM inference with multiple backends (LiteLLM, vLLM, llama.cpp)"
|
||||
),
|
||||
version="0.1.0",
|
||||
lifespan=lifespan,
|
||||
servers=[{"url": settings.external_url, "description": "LLM Inference 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-llm-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-llm-api instrumented -> {_otel_ep}")
|
||||
|
||||
except ImportError as _e:
|
||||
|
||||
print(f"[otel] skip: {_e}")
|
||||
|
||||
|
||||
# Runtime config client — instantiated here so middleware can capture it.
|
||||
# Started/stopped inside lifespan().
|
||||
runtime_config = RuntimeConfigClient(
|
||||
dashboard_url=settings.dashboard_url,
|
||||
live_log_logger_name="llm_inference",
|
||||
live_log_key="llm.log.level",
|
||||
)
|
||||
app.state.runtime_config = runtime_config
|
||||
|
||||
# Add middleware (order matters - first added is outermost)
|
||||
app.add_middleware(
|
||||
RateLimitMiddleware,
|
||||
rate=settings.rate_limit_rps,
|
||||
burst=settings.rate_limit_burst,
|
||||
exclude_paths=["/health", "/ready"],
|
||||
runtime_config=runtime_config,
|
||||
rate_key="llm.rate_limit.rps",
|
||||
burst_key="llm.rate_limit.burst",
|
||||
)
|
||||
|
||||
app.add_middleware(RequestIdMiddleware)
|
||||
|
||||
app.include_router(health.router, tags=["Health"])
|
||||
app.include_router(completions.router, prefix="/v1", tags=["Completions"])
|
||||
app.include_router(models.router, prefix="/v1", tags=["Models"])
|
||||
app.include_router(info.router, tags=["Catalog"])
|
||||
|
||||
return app
|
||||
|
|
@ -0,0 +1,232 @@
|
|||
"""FastAPI dependencies for LLM inference."""
|
||||
|
||||
import asyncio
|
||||
import hmac
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import Header, HTTPException, Request
|
||||
|
||||
from llm_inference.client import LLMClient
|
||||
from llm_inference.config import LLMSettings
|
||||
from llm_inference.logging import get_logger
|
||||
|
||||
logger = get_logger("dependencies")
|
||||
|
||||
|
||||
def get_client(request: Request) -> LLMClient:
|
||||
"""Get the LLM client from application state.
|
||||
|
||||
Args:
|
||||
request: FastAPI request object.
|
||||
|
||||
Returns:
|
||||
LLMClient: The LLM client instance.
|
||||
"""
|
||||
return request.app.state.client
|
||||
|
||||
|
||||
def get_settings(request: Request) -> LLMSettings:
|
||||
"""Get settings from application state.
|
||||
|
||||
Args:
|
||||
request: FastAPI request object.
|
||||
|
||||
Returns:
|
||||
LLMSettings: Application settings.
|
||||
"""
|
||||
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.
|
||||
|
||||
This dependency checks the Authorization header for a valid Bearer token.
|
||||
If authentication is disabled (no tokens configured), returns None.
|
||||
If authentication is enabled, validates the token and returns it.
|
||||
|
||||
Args:
|
||||
request: FastAPI request object.
|
||||
authorization: Authorization header value.
|
||||
|
||||
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
|
||||
|
||||
# Auth disabled - allow all requests
|
||||
if not settings.auth_enabled:
|
||||
return None
|
||||
|
||||
# Auth enabled - validate token
|
||||
if not authorization:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail={
|
||||
"error": "Authentication required",
|
||||
"message": "Missing Authorization header",
|
||||
},
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
|
||||
# Parse Bearer token
|
||||
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]
|
||||
|
||||
# Constant-time comparison against all valid tokens
|
||||
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 completions to prevent resource exhaustion.
|
||||
|
||||
Uses a semaphore to limit the number of concurrent completion requests.
|
||||
Returns 503 Service Unavailable when limit is exceeded.
|
||||
"""
|
||||
|
||||
def __init__(self, max_concurrent: int) -> None:
|
||||
"""Initialize concurrency limiter.
|
||||
|
||||
Args:
|
||||
max_concurrent: Maximum number of concurrent completions.
|
||||
"""
|
||||
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 completion slot.
|
||||
|
||||
Args:
|
||||
blocking: If True, wait for a slot. If False (default), return 503 immediately.
|
||||
|
||||
Yields:
|
||||
None when slot is acquired.
|
||||
|
||||
Raises:
|
||||
HTTPException: 503 if no slots available and blocking=False.
|
||||
"""
|
||||
if not blocking and self._semaphore.locked():
|
||||
# Non-blocking mode and no slots available - reject immediately
|
||||
# Note: small race window exists but semaphore still enforces limit
|
||||
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"},
|
||||
)
|
||||
|
||||
# Acquire semaphore (blocks if blocking=True and no slots, or immediate if available)
|
||||
await self._semaphore.acquire()
|
||||
|
||||
# Update counter with lock for thread safety
|
||||
async with self._counter_lock:
|
||||
self._current += 1
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
async with self._counter_lock:
|
||||
self._current -= 1
|
||||
self._semaphore.release()
|
||||
|
||||
|
||||
# Global concurrency limiter instance (initialized in app startup)
|
||||
_concurrency_limiter: ConcurrencyLimiter | None = None
|
||||
|
||||
|
||||
def init_concurrency_limiter(max_concurrent: int) -> ConcurrencyLimiter:
|
||||
"""Initialize the global concurrency limiter.
|
||||
|
||||
Args:
|
||||
max_concurrent: Maximum concurrent completions.
|
||||
|
||||
Returns:
|
||||
ConcurrencyLimiter: The initialized 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.
|
||||
|
||||
Returns:
|
||||
ConcurrencyLimiter: The global limiter instance.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If limiter not initialized.
|
||||
"""
|
||||
if _concurrency_limiter is None:
|
||||
raise RuntimeError("Concurrency limiter not initialized")
|
||||
return _concurrency_limiter
|
||||
|
||||
|
||||
async def require_completion_slot(request: Request) -> AsyncIterator[None]:
|
||||
"""FastAPI dependency that acquires a completion slot.
|
||||
|
||||
Args:
|
||||
request: FastAPI request.
|
||||
|
||||
Yields:
|
||||
None when slot is acquired.
|
||||
|
||||
Raises:
|
||||
HTTPException: 503 if no slots available.
|
||||
"""
|
||||
limiter = get_concurrency_limiter()
|
||||
async with limiter.acquire():
|
||||
yield
|
||||
|
|
@ -0,0 +1,230 @@
|
|||
"""FastAPI middleware for request handling and rate limiting.
|
||||
|
||||
Provides:
|
||||
- RequestIdMiddleware: Extracts/generates request IDs and propagates via context
|
||||
- RateLimitMiddleware: Token bucket rate limiting with 429 responses
|
||||
"""
|
||||
|
||||
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 llm_inference.logging import get_logger, set_request_id
|
||||
|
||||
logger = get_logger("middleware")
|
||||
|
||||
|
||||
class RequestIdMiddleware(BaseHTTPMiddleware):
|
||||
"""Middleware to extract or generate request IDs.
|
||||
|
||||
Extracts X-Request-ID from incoming headers or generates a new UUID.
|
||||
Sets the request ID in context for logging and returns it in response headers.
|
||||
"""
|
||||
|
||||
async def dispatch(
|
||||
self,
|
||||
request: Request,
|
||||
call_next: Callable[[Request], Awaitable[Response]],
|
||||
) -> Response:
|
||||
"""Process request with request ID tracking.
|
||||
|
||||
Args:
|
||||
request: Incoming HTTP request.
|
||||
call_next: Next middleware/handler in chain.
|
||||
|
||||
Returns:
|
||||
Response: HTTP response with X-Request-ID header.
|
||||
"""
|
||||
# Extract from header or generate new
|
||||
request_id = request.headers.get("X-Request-ID")
|
||||
if not request_id:
|
||||
request_id = str(uuid.uuid4())
|
||||
|
||||
# Set in context for logging
|
||||
set_request_id(request_id)
|
||||
|
||||
# Store in request state for handlers
|
||||
request.state.request_id = request_id
|
||||
|
||||
try:
|
||||
response = await call_next(request)
|
||||
response.headers["X-Request-ID"] = request_id
|
||||
return response
|
||||
finally:
|
||||
# Clear context after request
|
||||
set_request_id(None)
|
||||
|
||||
|
||||
class TokenBucket:
|
||||
"""Token bucket for rate limiting with optional live tuning.
|
||||
|
||||
Implements a simple token bucket algorithm where tokens are added
|
||||
at a fixed rate up to a maximum burst size. If a runtime_config and
|
||||
rate/burst keys are supplied, the bucket reads the latest override
|
||||
on every acquire() — making the limit live-tunable from dashboard
|
||||
without restart.
|
||||
"""
|
||||
|
||||
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. When set together
|
||||
with rate_key / burst_key, the bucket pulls live values on
|
||||
each acquire() and falls back to the constructor values.
|
||||
rate_key: Dashboard config key for live rate (e.g., "llm.rate_limit.rps").
|
||||
burst_key: Dashboard config key for live 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:
|
||||
# On burst increase, top up; on decrease, clamp tokens to new 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.
|
||||
|
||||
Returns:
|
||||
bool: True if token acquired, False if rate limited.
|
||||
"""
|
||||
self._refresh_from_config()
|
||||
now = time.monotonic()
|
||||
elapsed = now - self.last_update
|
||||
self.last_update = now
|
||||
|
||||
# Add tokens based on elapsed time
|
||||
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.
|
||||
|
||||
Returns:
|
||||
float: Seconds to wait before retrying.
|
||||
"""
|
||||
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.
|
||||
|
||||
Returns 429 Too Many Requests with Retry-After header when exceeded.
|
||||
|
||||
WARNING: This is per-process rate limiting. In multi-replica deployments
|
||||
(e.g., Kubernetes with multiple pods), each replica maintains its own
|
||||
independent rate limit. A configured limit of 10 RPS with 5 replicas
|
||||
effectively allows 50 RPS total.
|
||||
|
||||
For distributed rate limiting in production, use an external solution:
|
||||
- Redis-based rate limiting (e.g., redis-rate-limiter)
|
||||
- API gateway rate limiting (e.g., Kong, nginx)
|
||||
- Cloud provider rate limiting (e.g., AWS API Gateway)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
app: object,
|
||||
rate: float = 10.0,
|
||||
burst: int = 20,
|
||||
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.
|
||||
|
||||
Args:
|
||||
app: FastAPI application.
|
||||
rate: Requests per second limit (fallback when no runtime_config).
|
||||
burst: Maximum burst size (fallback).
|
||||
exclude_paths: Paths to exclude from rate limiting.
|
||||
runtime_config: Optional RuntimeConfigClient to enable live rate tuning.
|
||||
rate_key: Dashboard config key (e.g., "llm.rate_limit.rps").
|
||||
burst_key: Dashboard config key (e.g., "llm.rate_limit.burst").
|
||||
"""
|
||||
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.
|
||||
|
||||
Args:
|
||||
request: Incoming HTTP request.
|
||||
call_next: Next middleware/handler in chain.
|
||||
|
||||
Returns:
|
||||
Response: HTTP response or 429 if rate limited.
|
||||
"""
|
||||
# Skip rate limiting for excluded paths
|
||||
if request.url.path in self.exclude_paths:
|
||||
return await call_next(request)
|
||||
|
||||
# Try to acquire token
|
||||
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,5 @@
|
|||
"""API routes for LLM inference."""
|
||||
|
||||
from llm_inference.api.routes import completions, health, info, models
|
||||
|
||||
__all__ = ["completions", "health", "info", "models"]
|
||||
|
|
@ -0,0 +1,221 @@
|
|||
"""Chat completion routes with streaming support."""
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncGenerator
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from sse_starlette.sse import EventSourceResponse
|
||||
|
||||
from llm_inference.api.dependencies import (
|
||||
ConcurrencyLimiter,
|
||||
get_client,
|
||||
get_concurrency_limiter,
|
||||
verify_bearer_token,
|
||||
)
|
||||
from llm_inference.client import LLMClient
|
||||
from llm_inference.exceptions import (
|
||||
BackendNotAvailableError,
|
||||
BackendNotEnabledError,
|
||||
CompletionError,
|
||||
)
|
||||
from llm_inference.schemas import (
|
||||
CompletionRequest,
|
||||
CompletionResponse,
|
||||
TextCompletionRequest,
|
||||
TextCompletionResponse,
|
||||
)
|
||||
|
||||
router = APIRouter(dependencies=[Depends(verify_bearer_token)])
|
||||
|
||||
|
||||
@router.post(
|
||||
"/chat/completions",
|
||||
response_model=None, # Disable auto-generation due to Union with EventSourceResponse
|
||||
)
|
||||
async def chat_completions(
|
||||
request: CompletionRequest,
|
||||
client: LLMClient = Depends(get_client),
|
||||
) -> CompletionResponse | EventSourceResponse:
|
||||
"""OpenAI-compatible chat completions endpoint.
|
||||
|
||||
Supports both streaming (SSE) and non-streaming responses.
|
||||
Backend can be overridden per-request via the `backend` field.
|
||||
|
||||
Concurrency is limited per-process. For streaming requests, the slot is
|
||||
held for the entire duration of the stream.
|
||||
|
||||
Args:
|
||||
request: Completion request with messages, model, and options.
|
||||
client: LLM client instance.
|
||||
|
||||
Returns:
|
||||
CompletionResponse or EventSourceResponse for streaming.
|
||||
|
||||
Raises:
|
||||
HTTPException: If the request fails or concurrency limit exceeded (503).
|
||||
"""
|
||||
limiter = get_concurrency_limiter()
|
||||
|
||||
if request.stream:
|
||||
# For streaming: wrap generator to hold slot throughout entire stream
|
||||
# This is critical - the slot must be held until streaming completes
|
||||
return EventSourceResponse(
|
||||
_stream_with_slot(_stream_generator(client, request), limiter),
|
||||
media_type="text/event-stream",
|
||||
ping=15, # Send ping every 15s to detect dead connections
|
||||
)
|
||||
|
||||
# For non-streaming: acquire slot, run completion, release slot
|
||||
async with limiter.acquire():
|
||||
try:
|
||||
kwargs = _build_completion_kwargs(request)
|
||||
|
||||
response = await client.complete(
|
||||
messages=request.messages,
|
||||
model=request.model,
|
||||
backend=request.backend,
|
||||
**kwargs,
|
||||
)
|
||||
return response
|
||||
|
||||
except (BackendNotAvailableError, BackendNotEnabledError) as e:
|
||||
raise HTTPException(status_code=400, detail=str(e)) from None
|
||||
except CompletionError as e:
|
||||
raise HTTPException(status_code=500, detail=str(e)) from None
|
||||
|
||||
|
||||
@router.post("/completions", response_model=TextCompletionResponse)
|
||||
async def text_completions(
|
||||
request: TextCompletionRequest,
|
||||
client: LLMClient = Depends(get_client),
|
||||
) -> TextCompletionResponse:
|
||||
"""OpenAI-compatible legacy text completions endpoint (/v1/completions).
|
||||
|
||||
Non-streaming. Routes to the model's backend (vLLM primary, cloud via
|
||||
LiteLLM). Backends that do not support text completion return 501.
|
||||
"""
|
||||
limiter = get_concurrency_limiter()
|
||||
async with limiter.acquire():
|
||||
try:
|
||||
kwargs: dict[str, object] = {"temperature": request.temperature}
|
||||
if request.max_tokens is not None:
|
||||
kwargs["max_tokens"] = request.max_tokens
|
||||
if request.top_p is not None:
|
||||
kwargs["top_p"] = request.top_p
|
||||
if request.frequency_penalty is not None:
|
||||
kwargs["frequency_penalty"] = request.frequency_penalty
|
||||
if request.presence_penalty is not None:
|
||||
kwargs["presence_penalty"] = request.presence_penalty
|
||||
if request.stop is not None:
|
||||
stop_seq = (
|
||||
[request.stop]
|
||||
if isinstance(request.stop, str)
|
||||
else list(request.stop)
|
||||
)
|
||||
stop_seq = [s for s in stop_seq if s]
|
||||
if stop_seq:
|
||||
kwargs["stop"] = stop_seq
|
||||
|
||||
return await client.text_complete(
|
||||
prompt=request.prompt,
|
||||
model=request.model,
|
||||
backend=request.backend,
|
||||
**kwargs,
|
||||
)
|
||||
except NotImplementedError as e:
|
||||
raise HTTPException(status_code=501, detail=str(e)) from None
|
||||
except (BackendNotAvailableError, BackendNotEnabledError) as e:
|
||||
raise HTTPException(status_code=400, detail=str(e)) from None
|
||||
except CompletionError as e:
|
||||
raise HTTPException(status_code=500, detail=str(e)) from None
|
||||
|
||||
|
||||
async def _stream_generator(
|
||||
client: LLMClient,
|
||||
request: CompletionRequest,
|
||||
) -> AsyncGenerator[dict[str, str], None]:
|
||||
"""Generate SSE events for streaming response.
|
||||
|
||||
Args:
|
||||
client: LLM client instance.
|
||||
request: Completion request.
|
||||
|
||||
Yields:
|
||||
SSE event dictionaries with data field.
|
||||
"""
|
||||
try:
|
||||
kwargs = _build_completion_kwargs(request)
|
||||
|
||||
async for chunk in client.stream(
|
||||
messages=request.messages,
|
||||
model=request.model,
|
||||
backend=request.backend,
|
||||
**kwargs,
|
||||
):
|
||||
yield {"data": json.dumps(chunk.model_dump())}
|
||||
|
||||
# Send [DONE] marker to signal end of stream
|
||||
yield {"data": "[DONE]"}
|
||||
|
||||
except (BackendNotAvailableError, BackendNotEnabledError, CompletionError) as e:
|
||||
# Send error event followed by [DONE] to signal stream termination
|
||||
yield {
|
||||
"event": "error",
|
||||
"data": json.dumps({"error": str(e)}),
|
||||
}
|
||||
yield {"data": "[DONE]"}
|
||||
|
||||
|
||||
async def _stream_with_slot(
|
||||
generator: AsyncGenerator[dict[str, str], None],
|
||||
limiter: ConcurrencyLimiter,
|
||||
) -> AsyncGenerator[dict[str, str], None]:
|
||||
"""Wrap a stream generator to hold a concurrency slot throughout streaming.
|
||||
|
||||
This ensures the concurrency limiter slot is held for the entire duration
|
||||
of the SSE stream, not just until the EventSourceResponse is returned.
|
||||
|
||||
Args:
|
||||
generator: The underlying stream generator.
|
||||
limiter: Concurrency limiter to acquire slot from.
|
||||
|
||||
Yields:
|
||||
SSE event dictionaries from the wrapped generator.
|
||||
"""
|
||||
async with limiter.acquire():
|
||||
async for item in generator:
|
||||
yield item
|
||||
|
||||
|
||||
def _build_completion_kwargs(request: CompletionRequest) -> dict[str, object]:
|
||||
"""Build kwargs dict from completion request.
|
||||
|
||||
Args:
|
||||
request: Completion request.
|
||||
|
||||
Returns:
|
||||
dict: Keyword arguments for completion call.
|
||||
"""
|
||||
kwargs: dict[str, object] = {
|
||||
"temperature": request.temperature,
|
||||
}
|
||||
|
||||
if request.max_tokens is not None:
|
||||
kwargs["max_tokens"] = request.max_tokens
|
||||
if request.top_p is not None:
|
||||
kwargs["top_p"] = request.top_p
|
||||
if request.frequency_penalty is not None:
|
||||
kwargs["frequency_penalty"] = request.frequency_penalty
|
||||
if request.presence_penalty is not None:
|
||||
kwargs["presence_penalty"] = request.presence_penalty
|
||||
if request.stop is not None:
|
||||
# Normalize to list and filter out empty strings
|
||||
# Some backends choke on empty stop sequences
|
||||
stop_seq = (
|
||||
[request.stop] if isinstance(request.stop, str) else list(request.stop)
|
||||
)
|
||||
stop_seq = [s for s in stop_seq if s] # Filter empty strings
|
||||
if stop_seq:
|
||||
kwargs["stop"] = stop_seq
|
||||
|
||||
return kwargs
|
||||
|
|
@ -0,0 +1,78 @@
|
|||
"""Health check routes."""
|
||||
|
||||
from fastapi import APIRouter, Depends
|
||||
|
||||
from llm_inference.api.dependencies import get_client
|
||||
from llm_inference.client import LLMClient
|
||||
from llm_inference.logging import get_logger
|
||||
from llm_inference.schemas import BackendHealth, HealthResponse, ReadinessResponse
|
||||
|
||||
router = APIRouter()
|
||||
logger = get_logger("routes.health")
|
||||
|
||||
|
||||
@router.get("/health", response_model=HealthResponse)
|
||||
async def health_check(
|
||||
client: LLMClient = Depends(get_client),
|
||||
) -> HealthResponse:
|
||||
"""Health check endpoint.
|
||||
|
||||
Returns the overall health status and per-backend health status.
|
||||
|
||||
Returns:
|
||||
HealthResponse: Overall 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,
|
||||
)
|
||||
)
|
||||
|
||||
# Determine overall status
|
||||
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: LLMClient = Depends(get_client),
|
||||
) -> ReadinessResponse:
|
||||
"""Readiness probe for Kubernetes.
|
||||
|
||||
Verifies that the default backend is healthy and able to serve requests.
|
||||
This is the critical check for load balancer routing.
|
||||
|
||||
Returns:
|
||||
ReadinessResponse: Readiness status (ready if default backend is healthy).
|
||||
"""
|
||||
try:
|
||||
# Get the default backend and check its health
|
||||
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,468 @@
|
|||
"""Component information endpoint for service catalog."""
|
||||
|
||||
from fastapi import APIRouter, Depends
|
||||
|
||||
from llm_inference.api.dependencies import get_client
|
||||
from llm_inference.client import LLMClient
|
||||
from llm_inference.config import SettingsCache
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/v1/info")
|
||||
async def get_component_info(
|
||||
client: LLMClient = Depends(get_client),
|
||||
) -> dict:
|
||||
"""Get component information for service catalog.
|
||||
|
||||
Returns complete metadata about this service including:
|
||||
- Resource information (component metadata)
|
||||
- Available models (from all enabled backends)
|
||||
- Available functions (API endpoints)
|
||||
|
||||
This endpoint is used by the catalog-api to aggregate service information
|
||||
and by backend systems to populate the catalog database.
|
||||
|
||||
Returns:
|
||||
dict: Component information matching catalog.resources, catalog.models,
|
||||
and catalog.functions schemas.
|
||||
"""
|
||||
settings = SettingsCache.get()
|
||||
|
||||
# Determine which backends are available
|
||||
backends = ["litellm"] # Always available
|
||||
if settings.enable_vllm:
|
||||
backends.append("vllm")
|
||||
if settings.enable_llamacpp:
|
||||
backends.append("llamacpp")
|
||||
|
||||
# Build resource information (maps to catalog.resources)
|
||||
resource = {
|
||||
"name": "LLM Inference Gateway",
|
||||
"slug": "llm-inference",
|
||||
"resource_type": "api_service",
|
||||
"provider": "internal",
|
||||
"base_url": f"http://didiAI-llm-api:{settings.port}",
|
||||
"configuration": {
|
||||
"version": "1.0.0",
|
||||
"port": settings.port,
|
||||
"external_url": settings.external_url,
|
||||
"default_backend": settings.default_backend,
|
||||
"backends": backends,
|
||||
"default_model": settings.default_model,
|
||||
},
|
||||
"authentication": {
|
||||
"type": "bearer",
|
||||
"required": bool(settings.api_tokens),
|
||||
"env_var": "LLM_API_TOKENS",
|
||||
},
|
||||
"headers": {
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "application/json",
|
||||
},
|
||||
"rate_limits": {
|
||||
"requests_per_second": settings.rate_limit_rps,
|
||||
"burst": settings.rate_limit_burst,
|
||||
"concurrent": settings.max_concurrent_completions,
|
||||
},
|
||||
"cost_tracking": {
|
||||
"enabled": False,
|
||||
},
|
||||
"tags": ["llm", "inference", "openai-compatible", "gateway", "nlp"],
|
||||
"is_active": True,
|
||||
"metadata": {
|
||||
"category": "nlp",
|
||||
"gpu_required": settings.enable_vllm,
|
||||
"status": "healthy",
|
||||
"vllm_base_url": settings.vllm_base_url if settings.enable_vllm else None,
|
||||
"llamacpp_base_url": (
|
||||
settings.llamacpp_base_url if settings.enable_llamacpp else None
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
# Collect models from all backends
|
||||
models = []
|
||||
|
||||
# Get models from enabled backends
|
||||
try:
|
||||
all_models = await client.list_models()
|
||||
|
||||
for model_info in all_models:
|
||||
model_entry = {
|
||||
"name": model_info.id,
|
||||
"slug": model_info.id.lower().replace("/", "-").replace("_", "-"),
|
||||
"provider": _get_provider_from_model_id(model_info.id),
|
||||
"model_type": _get_model_type(model_info),
|
||||
"capabilities": model_info.capabilities or [],
|
||||
"configuration": {
|
||||
"backend": model_info.backend,
|
||||
"context_length": model_info.context_length,
|
||||
"loaded": model_info.loaded,
|
||||
},
|
||||
"endpoint": _get_model_endpoint(model_info, settings),
|
||||
"api_key_ref": None,
|
||||
"tags": _get_model_tags(model_info),
|
||||
"is_active": model_info.loaded,
|
||||
"metadata": {
|
||||
"backend": model_info.backend,
|
||||
"model_id": model_info.id,
|
||||
},
|
||||
}
|
||||
models.append(model_entry)
|
||||
except Exception:
|
||||
# If model listing fails, continue with empty models list
|
||||
pass
|
||||
|
||||
# Define available functions (maps to catalog.functions)
|
||||
functions = [
|
||||
{
|
||||
"name": "Chat Completions",
|
||||
"slug": "llm-chat-completions",
|
||||
"category": "completion",
|
||||
"description": (
|
||||
"Create a chat completion with optional streaming support. "
|
||||
"OpenAI-compatible endpoint supporting multiple backends."
|
||||
),
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"messages": {
|
||||
"type": "array",
|
||||
"description": "List of messages (1-1000)",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"role": {
|
||||
"type": "string",
|
||||
"enum": ["system", "user", "assistant", "function", "tool"],
|
||||
},
|
||||
"content": {"type": "string"},
|
||||
},
|
||||
"required": ["role"],
|
||||
},
|
||||
},
|
||||
"model": {"type": "string", "description": "Model identifier"},
|
||||
"temperature": {
|
||||
"type": "number",
|
||||
"minimum": 0.0,
|
||||
"maximum": 2.0,
|
||||
"default": 0.7,
|
||||
},
|
||||
"max_tokens": {"type": "integer", "minimum": 1, "maximum": 1000000},
|
||||
"stream": {"type": "boolean", "default": False},
|
||||
"backend": {
|
||||
"type": "string",
|
||||
"enum": backends,
|
||||
"description": "Backend override",
|
||||
},
|
||||
"top_p": {"type": "number", "minimum": 0.0, "maximum": 1.0},
|
||||
"frequency_penalty": {"type": "number", "minimum": -2.0, "maximum": 2.0},
|
||||
"presence_penalty": {"type": "number", "minimum": -2.0, "maximum": 2.0},
|
||||
"stop": {
|
||||
"oneOf": [
|
||||
{"type": "string"},
|
||||
{"type": "array", "items": {"type": "string"}},
|
||||
]
|
||||
},
|
||||
},
|
||||
"required": ["messages", "model"],
|
||||
},
|
||||
"output_schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"id": {"type": "string"},
|
||||
"object": {"type": "string", "const": "chat.completion"},
|
||||
"created": {"type": "integer"},
|
||||
"model": {"type": "string"},
|
||||
"choices": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"index": {"type": "integer"},
|
||||
"message": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"role": {"type": "string"},
|
||||
"content": {"type": "string"},
|
||||
},
|
||||
},
|
||||
"finish_reason": {"type": "string"},
|
||||
},
|
||||
},
|
||||
},
|
||||
"usage": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"prompt_tokens": {"type": "integer"},
|
||||
"completion_tokens": {"type": "integer"},
|
||||
"total_tokens": {"type": "integer"},
|
||||
},
|
||||
},
|
||||
"backend": {"type": "string"},
|
||||
},
|
||||
},
|
||||
"implementation": {
|
||||
"method": "POST",
|
||||
"path": "/v1/chat/completions",
|
||||
"content_type": "application/json",
|
||||
"timeout": 120,
|
||||
"supports_streaming": True,
|
||||
},
|
||||
"endpoint": f"http://localhost:{settings.port}/v1/chat/completions",
|
||||
"tags": ["llm", "chat", "openai", "streaming"],
|
||||
"is_active": True,
|
||||
"metadata": {
|
||||
"rate_limited": True,
|
||||
"auth_required": bool(settings.api_tokens),
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "List Models",
|
||||
"slug": "llm-list-models",
|
||||
"category": "discovery",
|
||||
"description": "List all available models across all enabled backends",
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"backend": {
|
||||
"type": "string",
|
||||
"enum": backends,
|
||||
"description": "Optional backend filter",
|
||||
}
|
||||
},
|
||||
},
|
||||
"output_schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"object": {"type": "string", "const": "list"},
|
||||
"data": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"id": {"type": "string"},
|
||||
"backend": {"type": "string"},
|
||||
"loaded": {"type": "boolean"},
|
||||
"context_length": {"type": "integer"},
|
||||
"capabilities": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"implementation": {
|
||||
"method": "GET",
|
||||
"path": "/v1/models",
|
||||
"timeout": 10,
|
||||
},
|
||||
"endpoint": f"http://localhost:{settings.port}/v1/models",
|
||||
"tags": ["discovery", "models"],
|
||||
"is_active": True,
|
||||
"metadata": {},
|
||||
},
|
||||
{
|
||||
"name": "List Backends",
|
||||
"slug": "llm-list-backends",
|
||||
"category": "discovery",
|
||||
"description": "List all available backends",
|
||||
"input_schema": {"type": "object", "properties": {}},
|
||||
"output_schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"backends": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
}
|
||||
},
|
||||
},
|
||||
"implementation": {
|
||||
"method": "GET",
|
||||
"path": "/v1/backends",
|
||||
"timeout": 5,
|
||||
},
|
||||
"endpoint": f"http://localhost:{settings.port}/v1/backends",
|
||||
"tags": ["discovery"],
|
||||
"is_active": True,
|
||||
"metadata": {},
|
||||
},
|
||||
{
|
||||
"name": "Load Model",
|
||||
"slug": "llm-load-model",
|
||||
"category": "management",
|
||||
"description": "Load a model on a local backend (vLLM or llama.cpp only)",
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"model": {"type": "string"},
|
||||
"backend": {"type": "string", "enum": ["vllm", "llamacpp"]},
|
||||
},
|
||||
"required": ["model", "backend"],
|
||||
},
|
||||
"output_schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"success": {"type": "boolean"},
|
||||
"model": {"type": "string"},
|
||||
"backend": {"type": "string"},
|
||||
"message": {"type": "string"},
|
||||
},
|
||||
},
|
||||
"implementation": {
|
||||
"method": "POST",
|
||||
"path": "/v1/models/load",
|
||||
"content_type": "application/json",
|
||||
"timeout": 60,
|
||||
},
|
||||
"endpoint": f"http://localhost:{settings.port}/v1/models/load",
|
||||
"tags": ["management", "models"],
|
||||
"is_active": settings.enable_vllm or settings.enable_llamacpp,
|
||||
"metadata": {},
|
||||
},
|
||||
{
|
||||
"name": "Unload Model",
|
||||
"slug": "llm-unload-model",
|
||||
"category": "management",
|
||||
"description": "Unload a model from a local backend (vLLM or llama.cpp only)",
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"model": {"type": "string"},
|
||||
"backend": {"type": "string", "enum": ["vllm", "llamacpp"]},
|
||||
},
|
||||
"required": ["model", "backend"],
|
||||
},
|
||||
"output_schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"success": {"type": "boolean"},
|
||||
"model": {"type": "string"},
|
||||
"backend": {"type": "string"},
|
||||
"message": {"type": "string"},
|
||||
},
|
||||
},
|
||||
"implementation": {
|
||||
"method": "POST",
|
||||
"path": "/v1/models/unload",
|
||||
"content_type": "application/json",
|
||||
"timeout": 30,
|
||||
},
|
||||
"endpoint": f"http://localhost:{settings.port}/v1/models/unload",
|
||||
"tags": ["management", "models"],
|
||||
"is_active": settings.enable_vllm or settings.enable_llamacpp,
|
||||
"metadata": {},
|
||||
},
|
||||
]
|
||||
|
||||
return {
|
||||
"resource": resource,
|
||||
"models": models,
|
||||
"functions": functions,
|
||||
}
|
||||
|
||||
|
||||
def _get_provider_from_model_id(model_id: str) -> str:
|
||||
"""Extract provider name from model ID.
|
||||
|
||||
Args:
|
||||
model_id: Model identifier (e.g., "openai/gpt-4", "meta-llama/Llama-2-7b")
|
||||
|
||||
Returns:
|
||||
str: Provider name (e.g., "openai", "meta-llama", "unknown")
|
||||
"""
|
||||
if "/" in model_id:
|
||||
return model_id.split("/")[0]
|
||||
if "gpt" in model_id.lower():
|
||||
return "openai"
|
||||
if "claude" in model_id.lower():
|
||||
return "anthropic"
|
||||
if "llama" in model_id.lower():
|
||||
return "meta"
|
||||
if "qwen" in model_id.lower():
|
||||
return "qwen"
|
||||
if "mistral" in model_id.lower():
|
||||
return "mistralai"
|
||||
return "unknown"
|
||||
|
||||
|
||||
def _get_model_type(model_info) -> str:
|
||||
"""Determine model type from capabilities.
|
||||
|
||||
Args:
|
||||
model_info: Model information object
|
||||
|
||||
Returns:
|
||||
str: Model type (llm, vision, audio, etc.)
|
||||
"""
|
||||
capabilities = model_info.capabilities or []
|
||||
if "vision" in capabilities or "multimodal" in capabilities:
|
||||
return "vision"
|
||||
if "audio" in capabilities:
|
||||
return "audio"
|
||||
if "embedding" in capabilities:
|
||||
return "embedding"
|
||||
return "llm"
|
||||
|
||||
|
||||
def _get_model_tags(model_info) -> list[str]:
|
||||
"""Generate tags for a model.
|
||||
|
||||
Args:
|
||||
model_info: Model information object
|
||||
|
||||
Returns:
|
||||
list[str]: List of tags
|
||||
"""
|
||||
tags = []
|
||||
|
||||
# Add model type tags
|
||||
if "gpt" in model_info.id.lower():
|
||||
tags.extend(["gpt", "openai"])
|
||||
if "claude" in model_info.id.lower():
|
||||
tags.extend(["claude", "anthropic"])
|
||||
if "llama" in model_info.id.lower():
|
||||
tags.extend(["llama", "meta"])
|
||||
if "qwen" in model_info.id.lower():
|
||||
tags.extend(["qwen"])
|
||||
if "mistral" in model_info.id.lower():
|
||||
tags.extend(["mistral"])
|
||||
|
||||
# Add capability tags
|
||||
capabilities = model_info.capabilities or []
|
||||
tags.extend(capabilities)
|
||||
|
||||
# Add backend tag
|
||||
tags.append(model_info.backend)
|
||||
|
||||
# Add size tags if detectable
|
||||
model_id_lower = model_info.id.lower()
|
||||
if "7b" in model_id_lower:
|
||||
tags.append("7b")
|
||||
elif "13b" in model_id_lower:
|
||||
tags.append("13b")
|
||||
elif "30b" in model_id_lower:
|
||||
tags.append("30b")
|
||||
elif "70b" in model_id_lower:
|
||||
tags.append("70b")
|
||||
elif "120b" in model_id_lower:
|
||||
tags.append("120b")
|
||||
|
||||
return list(set(tags)) # Remove duplicates
|
||||
|
||||
|
||||
def _get_model_endpoint(model_info, settings) -> str:
|
||||
"""Get the endpoint URL for a model.
|
||||
|
||||
Args:
|
||||
model_info: Model information object
|
||||
settings: Application settings
|
||||
|
||||
Returns:
|
||||
str: Endpoint URL
|
||||
"""
|
||||
# Default to gateway endpoint
|
||||
return f"http://localhost:{settings.port}/v1/chat/completions"
|
||||
|
|
@ -0,0 +1,134 @@
|
|||
"""Model management routes."""
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
|
||||
from llm_inference.api.dependencies import get_client, verify_bearer_token
|
||||
from llm_inference.client import LLMClient
|
||||
from llm_inference.exceptions import (
|
||||
BackendNotAvailableError,
|
||||
BackendNotEnabledError,
|
||||
ModelLoadError,
|
||||
)
|
||||
from llm_inference.schemas import (
|
||||
BackendListResponse,
|
||||
ModelListResponse,
|
||||
ModelLoadRequest,
|
||||
ModelLoadResponse,
|
||||
)
|
||||
from llm_inference.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: LLMClient = Depends(get_client),
|
||||
) -> ModelListResponse:
|
||||
"""List available models.
|
||||
|
||||
Args:
|
||||
backend: Optional backend filter. If not specified, returns models
|
||||
from all available backends.
|
||||
client: LLM 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
|
||||
|
||||
|
||||
@router.post("/models/load", response_model=ModelLoadResponse)
|
||||
async def load_model(
|
||||
request: ModelLoadRequest,
|
||||
client: LLMClient = Depends(get_client),
|
||||
) -> ModelLoadResponse:
|
||||
"""Load a model on the specified backend.
|
||||
|
||||
This is only supported by local backends (vLLM, llama.cpp).
|
||||
|
||||
Args:
|
||||
request: Model load request.
|
||||
client: LLM client instance.
|
||||
|
||||
Returns:
|
||||
ModelLoadResponse: Result of the load operation.
|
||||
|
||||
Raises:
|
||||
HTTPException: If loading fails or backend doesn't support it.
|
||||
"""
|
||||
try:
|
||||
success = await client.load_model(request.model, request.backend)
|
||||
return ModelLoadResponse(
|
||||
success=success,
|
||||
model=request.model,
|
||||
backend=request.backend.value,
|
||||
message="Model loaded successfully" if success else "Model load failed",
|
||||
)
|
||||
except NotImplementedError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e)) from None
|
||||
except (BackendNotAvailableError, BackendNotEnabledError) as e:
|
||||
raise HTTPException(status_code=400, detail=str(e)) from None
|
||||
except ModelLoadError as e:
|
||||
raise HTTPException(status_code=500, detail=str(e)) from None
|
||||
|
||||
|
||||
@router.post("/models/unload", response_model=ModelLoadResponse)
|
||||
async def unload_model(
|
||||
request: ModelLoadRequest,
|
||||
client: LLMClient = Depends(get_client),
|
||||
) -> ModelLoadResponse:
|
||||
"""Unload a model from the specified backend.
|
||||
|
||||
This is only supported by local backends (vLLM, llama.cpp).
|
||||
|
||||
Args:
|
||||
request: Model unload request.
|
||||
client: LLM client instance.
|
||||
|
||||
Returns:
|
||||
ModelLoadResponse: Result of the unload operation.
|
||||
|
||||
Raises:
|
||||
HTTPException: If unloading fails or backend doesn't support it.
|
||||
"""
|
||||
try:
|
||||
success = await client.unload_model(request.model, request.backend)
|
||||
return ModelLoadResponse(
|
||||
success=success,
|
||||
model=request.model,
|
||||
backend=request.backend.value,
|
||||
message="Model unloaded successfully" if success else "Model unload failed",
|
||||
)
|
||||
except NotImplementedError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e)) from None
|
||||
except (BackendNotAvailableError, BackendNotEnabledError) as e:
|
||||
raise HTTPException(status_code=400, detail=str(e)) from None
|
||||
|
||||
|
||||
@router.get("/backends", response_model=BackendListResponse)
|
||||
async def list_backends(
|
||||
client: LLMClient = Depends(get_client),
|
||||
) -> BackendListResponse:
|
||||
"""List available backends.
|
||||
|
||||
Args:
|
||||
client: LLM 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 LLM inference."""
|
||||
|
||||
from llm_inference.backends.base import LLMBackend
|
||||
from llm_inference.backends.registry import BackendRegistry
|
||||
|
||||
__all__ = [
|
||||
"BackendRegistry",
|
||||
"LLMBackend",
|
||||
]
|
||||
|
|
@ -0,0 +1,146 @@
|
|||
"""Abstract base class for LLM backends."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
from llm_inference.schemas import (
|
||||
CompletionChunk,
|
||||
CompletionResponse,
|
||||
TextCompletionResponse,
|
||||
)
|
||||
from llm_inference.types import ChatMessage, ModelInfo
|
||||
|
||||
|
||||
class LLMBackend(ABC):
|
||||
"""Abstract base class for LLM 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., 'litellm', 'vllm').
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def complete(
|
||||
self,
|
||||
messages: list[ChatMessage],
|
||||
model: str,
|
||||
**kwargs: object,
|
||||
) -> CompletionResponse:
|
||||
"""Generate a chat completion.
|
||||
|
||||
Args:
|
||||
messages: List of chat messages in the conversation.
|
||||
model: Model identifier to use for completion.
|
||||
**kwargs: Additional parameters passed to the backend.
|
||||
|
||||
Returns:
|
||||
CompletionResponse: The completion response.
|
||||
|
||||
Raises:
|
||||
CompletionError: If the completion request fails.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def stream(
|
||||
self,
|
||||
messages: list[ChatMessage],
|
||||
model: str,
|
||||
**kwargs: object,
|
||||
) -> AsyncIterator[CompletionChunk]:
|
||||
"""Generate a streaming chat completion.
|
||||
|
||||
Args:
|
||||
messages: List of chat messages in the conversation.
|
||||
model: Model identifier to use for completion.
|
||||
**kwargs: Additional parameters passed to the backend.
|
||||
|
||||
Yields:
|
||||
CompletionChunk: Streaming chunks of the completion.
|
||||
|
||||
Raises:
|
||||
CompletionError: If the completion request fails.
|
||||
"""
|
||||
...
|
||||
yield # pragma: no cover (makes this a generator)
|
||||
|
||||
async def text_complete(
|
||||
self,
|
||||
prompt: str | list[str],
|
||||
model: str,
|
||||
**kwargs: object,
|
||||
) -> TextCompletionResponse:
|
||||
"""Generate a (legacy) text completion for ``prompt``.
|
||||
|
||||
Concrete by default so backends that cannot serve /v1/completions are
|
||||
not forced to implement it; override in OpenAI-compatible backends.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If the backend does not support text completion.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
f"Backend '{self.name}' does not support text completions"
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
async def list_models(self) -> list[ModelInfo]:
|
||||
"""List available models for this backend.
|
||||
|
||||
Returns:
|
||||
list[ModelInfo]: List of available models.
|
||||
"""
|
||||
...
|
||||
|
||||
async def load_model(self, model: str) -> bool:
|
||||
"""Load a model on this backend.
|
||||
|
||||
This is optional and only supported by local backends (vLLM, llama.cpp).
|
||||
|
||||
Args:
|
||||
model: Model identifier to load.
|
||||
|
||||
Returns:
|
||||
bool: True if the model was loaded successfully.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If the backend doesn't support model loading.
|
||||
ModelLoadError: If loading the model fails.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
f"Backend '{self.name}' does not support dynamic model loading"
|
||||
)
|
||||
|
||||
async def unload_model(self, model: str) -> bool:
|
||||
"""Unload a model from this backend.
|
||||
|
||||
This is optional and only supported by local backends (vLLM, llama.cpp).
|
||||
|
||||
Args:
|
||||
model: Model identifier to unload.
|
||||
|
||||
Returns:
|
||||
bool: True if the model was unloaded successfully.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If the backend doesn't support model unloading.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
f"Backend '{self.name}' does not support dynamic model unloading"
|
||||
)
|
||||
|
||||
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,518 @@
|
|||
"""LiteLLM backend implementation.
|
||||
|
||||
LiteLLM provides a unified interface for 100+ LLM providers including:
|
||||
- OpenAI
|
||||
- Anthropic
|
||||
- OpenRouter
|
||||
- Azure OpenAI
|
||||
- Google AI (Gemini)
|
||||
- AWS Bedrock
|
||||
- And many more
|
||||
"""
|
||||
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
import httpx
|
||||
import litellm
|
||||
import litellm.exceptions
|
||||
from litellm import acompletion, atext_completion
|
||||
|
||||
from llm_inference.backends.base import LLMBackend
|
||||
from llm_inference.config import LLMSettings
|
||||
from llm_inference.exceptions import (
|
||||
CompletionError,
|
||||
LLMRateLimitError,
|
||||
LLMTimeoutError,
|
||||
)
|
||||
from llm_inference.retry import retry_with_backoff
|
||||
from llm_inference.schemas import (
|
||||
CompletionChunk,
|
||||
CompletionResponse,
|
||||
TextChoice,
|
||||
TextCompletionResponse,
|
||||
)
|
||||
from llm_inference.types import (
|
||||
ChatMessage,
|
||||
Choice,
|
||||
Delta,
|
||||
ModelInfo,
|
||||
StreamChoice,
|
||||
Usage,
|
||||
)
|
||||
from llm_inference.utils import safe_close_stream
|
||||
|
||||
logger = logging.getLogger("llm_inference.backends.litellm")
|
||||
|
||||
|
||||
class LiteLLMBackend(LLMBackend):
|
||||
"""LiteLLM backend supporting 100+ LLM providers.
|
||||
|
||||
This is the default backend and handles most cloud-based LLM providers
|
||||
including OpenRouter for accessing various models through a single API.
|
||||
"""
|
||||
|
||||
# Default cache TTL for model list (1 hour)
|
||||
_MODELS_CACHE_TTL: float = 3600.0
|
||||
|
||||
def __init__(self, settings: LLMSettings) -> None:
|
||||
"""Initialize LiteLLM backend.
|
||||
|
||||
Args:
|
||||
settings: Application settings with API keys.
|
||||
"""
|
||||
self._settings = settings
|
||||
self._configure_litellm()
|
||||
# Model list cache: (timestamp, models)
|
||||
self._models_cache: tuple[float, list[ModelInfo]] | None = None
|
||||
|
||||
def _configure_litellm(self) -> None:
|
||||
"""Configure LiteLLM with API keys from settings."""
|
||||
# Set API keys if provided
|
||||
if self._settings.openrouter_api_key:
|
||||
litellm.openrouter_api_key = self._settings.openrouter_api_key
|
||||
if self._settings.openai_api_key:
|
||||
litellm.openai_api_key = self._settings.openai_api_key
|
||||
if self._settings.anthropic_api_key:
|
||||
litellm.anthropic_api_key = self._settings.anthropic_api_key
|
||||
|
||||
# Suppress verbose logging
|
||||
litellm.suppress_debug_info = True
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
"""Backend identifier name."""
|
||||
return "litellm"
|
||||
|
||||
async def complete(
|
||||
self,
|
||||
messages: list[ChatMessage],
|
||||
model: str,
|
||||
**kwargs: object,
|
||||
) -> CompletionResponse:
|
||||
"""Generate a chat completion via LiteLLM.
|
||||
|
||||
Args:
|
||||
messages: List of chat messages.
|
||||
model: Model identifier (e.g., 'gpt-4', 'claude-3-opus').
|
||||
**kwargs: Additional parameters (temperature, max_tokens, etc.).
|
||||
|
||||
Returns:
|
||||
CompletionResponse: The completion response.
|
||||
|
||||
Raises:
|
||||
CompletionError: If the request fails after retries.
|
||||
LLMRateLimitError: If rate limited.
|
||||
LLMTimeoutError: If the request times out.
|
||||
"""
|
||||
|
||||
async def _do_complete() -> CompletionResponse:
|
||||
response = await acompletion(
|
||||
model=model,
|
||||
messages=[m.model_dump(exclude_none=True) for m in messages],
|
||||
timeout=self._settings.request_timeout,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Convert response to our schema
|
||||
choices = []
|
||||
for choice in response.choices:
|
||||
choices.append(
|
||||
Choice(
|
||||
index=choice.index,
|
||||
message=ChatMessage(
|
||||
role=choice.message.role,
|
||||
content=choice.message.content,
|
||||
),
|
||||
finish_reason=choice.finish_reason,
|
||||
)
|
||||
)
|
||||
|
||||
usage = None
|
||||
if response.usage:
|
||||
usage = Usage(
|
||||
prompt_tokens=response.usage.prompt_tokens,
|
||||
completion_tokens=response.usage.completion_tokens,
|
||||
total_tokens=response.usage.total_tokens,
|
||||
)
|
||||
|
||||
return CompletionResponse(
|
||||
id=response.id,
|
||||
created=response.created,
|
||||
model=response.model,
|
||||
choices=choices,
|
||||
usage=usage,
|
||||
backend=self.name,
|
||||
)
|
||||
|
||||
return await retry_with_backoff(
|
||||
_do_complete,
|
||||
max_retries=self._settings.max_retries,
|
||||
min_wait=self._settings.retry_min_wait,
|
||||
max_wait=self._settings.retry_max_wait,
|
||||
backend=self.name,
|
||||
model=model,
|
||||
)
|
||||
|
||||
async def text_complete(
|
||||
self,
|
||||
prompt: str | list[str],
|
||||
model: str,
|
||||
**kwargs: object,
|
||||
) -> TextCompletionResponse:
|
||||
"""Generate a legacy text completion via LiteLLM (cloud providers)."""
|
||||
|
||||
async def _do_text() -> TextCompletionResponse:
|
||||
response = await atext_completion(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
timeout=self._settings.request_timeout,
|
||||
**kwargs,
|
||||
)
|
||||
choices = [
|
||||
TextChoice(
|
||||
index=getattr(c, "index", i),
|
||||
text=getattr(c, "text", ""),
|
||||
finish_reason=getattr(c, "finish_reason", None),
|
||||
)
|
||||
for i, c in enumerate(response.choices)
|
||||
]
|
||||
usage = None
|
||||
if response.usage:
|
||||
usage = Usage(
|
||||
prompt_tokens=response.usage.prompt_tokens,
|
||||
completion_tokens=response.usage.completion_tokens,
|
||||
total_tokens=response.usage.total_tokens,
|
||||
)
|
||||
return TextCompletionResponse(
|
||||
id=response.id,
|
||||
created=response.created,
|
||||
model=response.model,
|
||||
choices=choices,
|
||||
usage=usage,
|
||||
backend=self.name,
|
||||
)
|
||||
|
||||
return await retry_with_backoff(
|
||||
_do_text,
|
||||
max_retries=self._settings.max_retries,
|
||||
min_wait=self._settings.retry_min_wait,
|
||||
max_wait=self._settings.retry_max_wait,
|
||||
backend=self.name,
|
||||
model=model,
|
||||
)
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
messages: list[ChatMessage],
|
||||
model: str,
|
||||
**kwargs: object,
|
||||
) -> AsyncIterator[CompletionChunk]:
|
||||
"""Generate a streaming chat completion via LiteLLM.
|
||||
|
||||
Args:
|
||||
messages: List of chat messages.
|
||||
model: Model identifier.
|
||||
**kwargs: Additional parameters.
|
||||
|
||||
Yields:
|
||||
CompletionChunk: Streaming chunks of the completion.
|
||||
|
||||
Raises:
|
||||
CompletionError: If the request fails.
|
||||
LLMRateLimitError: If rate limited.
|
||||
LLMTimeoutError: If the request times out.
|
||||
"""
|
||||
response = None
|
||||
try:
|
||||
response = await acompletion(
|
||||
model=model,
|
||||
messages=[m.model_dump(exclude_none=True) for m in messages],
|
||||
stream=True,
|
||||
timeout=self._settings.request_timeout,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
async for chunk in response:
|
||||
choices = []
|
||||
if chunk.choices:
|
||||
for choice in chunk.choices:
|
||||
delta = Delta(
|
||||
role=getattr(choice.delta, "role", None),
|
||||
content=getattr(choice.delta, "content", None),
|
||||
)
|
||||
choices.append(
|
||||
StreamChoice(
|
||||
index=choice.index,
|
||||
delta=delta,
|
||||
finish_reason=choice.finish_reason,
|
||||
)
|
||||
)
|
||||
|
||||
yield CompletionChunk(
|
||||
id=chunk.id,
|
||||
created=int(time.time()),
|
||||
model=chunk.model or model,
|
||||
choices=choices,
|
||||
)
|
||||
|
||||
except litellm.exceptions.RateLimitError as e:
|
||||
logger.warning("Rate limited during stream: %s", str(e))
|
||||
raise LLMRateLimitError(str(e), backend=self.name) from e
|
||||
except litellm.exceptions.Timeout as e:
|
||||
logger.warning("Timeout during stream: %s", str(e))
|
||||
raise LLMTimeoutError(str(e), backend=self.name) from e
|
||||
except litellm.exceptions.AuthenticationError as e:
|
||||
logger.error("Authentication failed: %s", str(e))
|
||||
raise CompletionError(
|
||||
f"Authentication failed: {e}", backend=self.name, model=model
|
||||
) from e
|
||||
except litellm.exceptions.APIError as e:
|
||||
logger.error("API error during stream: %s", str(e))
|
||||
raise CompletionError(str(e), backend=self.name, model=model) from e
|
||||
except Exception as e:
|
||||
logger.error("Unexpected error during stream: %s", str(e))
|
||||
raise CompletionError(str(e), backend=self.name, model=model) from e
|
||||
finally:
|
||||
# Ensure stream is properly closed on cancellation or error
|
||||
await safe_close_stream(response, logger)
|
||||
|
||||
async def list_models(self) -> list[ModelInfo]:
|
||||
"""List available models via LiteLLM.
|
||||
|
||||
Uses a TTL-based cache to avoid hitting external APIs on every call.
|
||||
Falls back to a curated list if dynamic fetching fails.
|
||||
|
||||
Returns:
|
||||
list[ModelInfo]: List of available models.
|
||||
"""
|
||||
# Check cache first
|
||||
now = time.time()
|
||||
if self._models_cache is not None:
|
||||
cached_time, cached_models = self._models_cache
|
||||
if now - cached_time < self._MODELS_CACHE_TTL:
|
||||
logger.debug(
|
||||
"Returning cached model list (%d models)", len(cached_models)
|
||||
)
|
||||
return cached_models
|
||||
|
||||
models: list[ModelInfo] = []
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(self._settings.connect_timeout)
|
||||
) as client:
|
||||
# Try to fetch OpenAI models dynamically
|
||||
if self._settings.openai_api_key:
|
||||
try:
|
||||
response = await client.get(
|
||||
"https://api.openai.com/v1/models",
|
||||
headers={
|
||||
"Authorization": f"Bearer {self._settings.openai_api_key}"
|
||||
},
|
||||
)
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
for m in data.get("data", []):
|
||||
model_id = m.get("id", "")
|
||||
# Filter to chat models (gpt-*)
|
||||
if model_id.startswith("gpt-"):
|
||||
models.append(
|
||||
ModelInfo(
|
||||
id=model_id,
|
||||
backend=self.name,
|
||||
capabilities=["chat"],
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug("Failed to fetch OpenAI models dynamically: %s", e)
|
||||
# Fall back to curated list for OpenAI
|
||||
models.extend(self._get_fallback_openai_models())
|
||||
|
||||
# Try to fetch OpenRouter models dynamically
|
||||
if self._settings.openrouter_api_key:
|
||||
try:
|
||||
response = await client.get(
|
||||
"https://openrouter.ai/api/v1/models",
|
||||
headers={
|
||||
"Authorization": f"Bearer {self._settings.openrouter_api_key}"
|
||||
},
|
||||
)
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
for m in data.get("data", []):
|
||||
model_id = m.get("id", "")
|
||||
if model_id:
|
||||
context_length = m.get("context_length")
|
||||
models.append(
|
||||
ModelInfo(
|
||||
id=f"openrouter/{model_id}",
|
||||
backend=self.name,
|
||||
context_length=context_length,
|
||||
capabilities=["chat"],
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug("Failed to fetch OpenRouter models dynamically: %s", e)
|
||||
# Fall back to curated list for OpenRouter
|
||||
models.extend(self._get_fallback_openrouter_models())
|
||||
|
||||
# Anthropic doesn't have a public models endpoint, use curated list
|
||||
if self._settings.anthropic_api_key:
|
||||
models.extend(self._get_fallback_anthropic_models())
|
||||
|
||||
# Cache the result
|
||||
self._models_cache = (time.time(), models)
|
||||
logger.debug("Cached model list (%d models)", len(models))
|
||||
|
||||
return models
|
||||
|
||||
def _get_fallback_openai_models(self) -> list[ModelInfo]:
|
||||
"""Get fallback list of OpenAI models."""
|
||||
return [
|
||||
ModelInfo(
|
||||
id="gpt-4o",
|
||||
backend=self.name,
|
||||
context_length=128000,
|
||||
capabilities=["chat"],
|
||||
),
|
||||
ModelInfo(
|
||||
id="gpt-4o-mini",
|
||||
backend=self.name,
|
||||
context_length=128000,
|
||||
capabilities=["chat"],
|
||||
),
|
||||
ModelInfo(
|
||||
id="gpt-4-turbo",
|
||||
backend=self.name,
|
||||
context_length=128000,
|
||||
capabilities=["chat"],
|
||||
),
|
||||
ModelInfo(
|
||||
id="gpt-3.5-turbo",
|
||||
backend=self.name,
|
||||
context_length=16385,
|
||||
capabilities=["chat"],
|
||||
),
|
||||
]
|
||||
|
||||
def _get_fallback_anthropic_models(self) -> list[ModelInfo]:
|
||||
"""Get fallback list of Anthropic models."""
|
||||
return [
|
||||
ModelInfo(
|
||||
id="claude-3-5-sonnet-20241022",
|
||||
backend=self.name,
|
||||
context_length=200000,
|
||||
capabilities=["chat"],
|
||||
),
|
||||
ModelInfo(
|
||||
id="claude-3-opus-20240229",
|
||||
backend=self.name,
|
||||
context_length=200000,
|
||||
capabilities=["chat"],
|
||||
),
|
||||
ModelInfo(
|
||||
id="claude-3-haiku-20240307",
|
||||
backend=self.name,
|
||||
context_length=200000,
|
||||
capabilities=["chat"],
|
||||
),
|
||||
]
|
||||
|
||||
def _get_fallback_openrouter_models(self) -> list[ModelInfo]:
|
||||
"""Get fallback list of OpenRouter models."""
|
||||
return [
|
||||
ModelInfo(
|
||||
id="openrouter/anthropic/claude-3.5-sonnet",
|
||||
backend=self.name,
|
||||
capabilities=["chat"],
|
||||
),
|
||||
ModelInfo(
|
||||
id="openrouter/openai/gpt-4o",
|
||||
backend=self.name,
|
||||
capabilities=["chat"],
|
||||
),
|
||||
ModelInfo(
|
||||
id="openrouter/meta-llama/llama-3.1-405b-instruct",
|
||||
backend=self.name,
|
||||
capabilities=["chat"],
|
||||
),
|
||||
ModelInfo(
|
||||
id="openrouter/google/gemini-pro-1.5",
|
||||
backend=self.name,
|
||||
capabilities=["chat"],
|
||||
),
|
||||
]
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check if LiteLLM backend can reach at least one provider.
|
||||
|
||||
Actually tests connectivity to configured providers instead of
|
||||
always returning True.
|
||||
|
||||
Returns:
|
||||
bool: True if at least one provider is reachable.
|
||||
"""
|
||||
async with httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(self._settings.connect_timeout)
|
||||
) as client:
|
||||
# Check OpenAI
|
||||
if self._settings.openai_api_key:
|
||||
try:
|
||||
response = await client.get(
|
||||
"https://api.openai.com/v1/models",
|
||||
headers={
|
||||
"Authorization": f"Bearer {self._settings.openai_api_key}"
|
||||
},
|
||||
)
|
||||
# 200 = success, 401 = reachable but bad key (still healthy)
|
||||
# Treating 401 as healthy supports key rotation scenarios
|
||||
if response.status_code in (200, 401):
|
||||
if response.status_code == 401:
|
||||
logger.warning(
|
||||
"OpenAI: API key invalid but endpoint reachable"
|
||||
)
|
||||
return True
|
||||
except httpx.RequestError as e:
|
||||
logger.debug("OpenAI health check failed: %s", e)
|
||||
|
||||
# Check Anthropic
|
||||
if self._settings.anthropic_api_key:
|
||||
try:
|
||||
response = await client.get(
|
||||
"https://api.anthropic.com/v1/models",
|
||||
headers={
|
||||
"x-api-key": self._settings.anthropic_api_key,
|
||||
"anthropic-version": "2023-06-01",
|
||||
},
|
||||
)
|
||||
if response.status_code in (200, 401):
|
||||
if response.status_code == 401:
|
||||
logger.warning(
|
||||
"Anthropic: API key invalid but endpoint reachable"
|
||||
)
|
||||
return True
|
||||
except httpx.RequestError as e:
|
||||
logger.debug("Anthropic health check failed: %s", e)
|
||||
|
||||
# Check OpenRouter
|
||||
if self._settings.openrouter_api_key:
|
||||
try:
|
||||
response = await client.get(
|
||||
"https://openrouter.ai/api/v1/models",
|
||||
headers={
|
||||
"Authorization": f"Bearer {self._settings.openrouter_api_key}"
|
||||
},
|
||||
)
|
||||
if response.status_code in (200, 401):
|
||||
if response.status_code == 401:
|
||||
logger.warning(
|
||||
"OpenRouter: API key invalid but endpoint reachable"
|
||||
)
|
||||
return True
|
||||
except httpx.RequestError as e:
|
||||
logger.debug("OpenRouter health check failed: %s", e)
|
||||
|
||||
# No providers reachable or configured
|
||||
logger.warning("No LLM providers reachable")
|
||||
return False
|
||||
|
|
@ -0,0 +1,477 @@
|
|||
"""llama.cpp backend with load balancing across multiple servers.
|
||||
|
||||
Supports round-robin distribution with automatic failover and periodic
|
||||
health checks. When a server goes down, traffic is routed to healthy
|
||||
servers. When it recovers, it re-enters the rotation.
|
||||
|
||||
Configuration:
|
||||
Single server (legacy):
|
||||
LLM_LLAMACPP_BASE_URL=http://localhost:8102
|
||||
|
||||
Multiple servers (load balanced):
|
||||
LLM_LLAMACPP_BASE_URLS=http://10.11.10.18:14001,http://10.11.10.19:14001
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import httpx
|
||||
|
||||
from llm_inference.backends.base import LLMBackend
|
||||
from llm_inference.config import LLMSettings
|
||||
from llm_inference.exceptions import (
|
||||
CompletionError,
|
||||
LLMConnectionError,
|
||||
LLMRateLimitError,
|
||||
LLMTimeoutError,
|
||||
ModelListError,
|
||||
)
|
||||
from llm_inference.retry import retry_with_backoff
|
||||
from llm_inference.schemas import CompletionChunk, CompletionResponse
|
||||
from llm_inference.types import (
|
||||
ChatMessage,
|
||||
Choice,
|
||||
Delta,
|
||||
ModelInfo,
|
||||
StreamChoice,
|
||||
Usage,
|
||||
)
|
||||
from llm_inference.utils import safe_close_stream
|
||||
|
||||
try:
|
||||
import openai
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
OPENAI_AVAILABLE = True
|
||||
except ImportError:
|
||||
OPENAI_AVAILABLE = False
|
||||
|
||||
logger = logging.getLogger("llm_inference.backends.llamacpp")
|
||||
|
||||
|
||||
@dataclass
|
||||
class LlamaCppServer:
|
||||
"""State for a single llama.cpp server."""
|
||||
|
||||
url: str
|
||||
healthy: bool = True
|
||||
request_count: int = 0
|
||||
error_count: int = 0
|
||||
last_check: float = 0
|
||||
response_times: list[float] = field(default_factory=list)
|
||||
|
||||
@property
|
||||
def avg_response_ms(self) -> float:
|
||||
if not self.response_times:
|
||||
return 0
|
||||
recent = self.response_times[-50:]
|
||||
return sum(recent) / len(recent)
|
||||
|
||||
@property
|
||||
def short_name(self) -> str:
|
||||
"""Extract host:port for logging."""
|
||||
return self.url.replace("http://", "").replace("https://", "")
|
||||
|
||||
|
||||
class LlamaCppBackend(LLMBackend):
|
||||
"""llama.cpp backend with round-robin load balancing.
|
||||
|
||||
When multiple URLs are configured via LLM_LLAMACPP_BASE_URLS,
|
||||
requests are distributed round-robin across healthy servers.
|
||||
Unhealthy servers are skipped and periodically re-checked.
|
||||
"""
|
||||
|
||||
def __init__(self, settings: LLMSettings) -> None:
|
||||
if not OPENAI_AVAILABLE:
|
||||
raise ImportError(
|
||||
"openai package is required for llama.cpp backend. "
|
||||
"Install with: pip install llm-inference[llamacpp]"
|
||||
)
|
||||
|
||||
self._settings = settings
|
||||
self._timeout = httpx.Timeout(
|
||||
connect=settings.connect_timeout,
|
||||
read=settings.request_timeout,
|
||||
write=settings.request_timeout,
|
||||
pool=settings.connect_timeout,
|
||||
)
|
||||
|
||||
# Initialize servers and clients
|
||||
self._servers: list[LlamaCppServer] = []
|
||||
self._clients: dict[str, AsyncOpenAI] = {}
|
||||
for url in settings.llamacpp_urls:
|
||||
server = LlamaCppServer(url=url)
|
||||
self._servers.append(server)
|
||||
self._clients[url] = AsyncOpenAI(
|
||||
base_url=f"{url}/v1",
|
||||
api_key="not-needed",
|
||||
timeout=self._timeout,
|
||||
max_retries=0,
|
||||
)
|
||||
|
||||
self._current_index = 0
|
||||
self._lock = asyncio.Lock()
|
||||
self._health_task: asyncio.Task[None] | None = None
|
||||
self._health_check_interval = settings.llamacpp_health_check_interval
|
||||
|
||||
server_names = [s.short_name for s in self._servers]
|
||||
logger.info(
|
||||
"llama.cpp backend: %d server(s): %s", len(self._servers), server_names
|
||||
)
|
||||
|
||||
# Start health check loop if multiple servers
|
||||
if len(self._servers) > 1:
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
self._health_task = loop.create_task(
|
||||
self._periodic_health_check(self._health_check_interval)
|
||||
)
|
||||
except RuntimeError:
|
||||
# No running loop yet (e.g., during import). Will be started later.
|
||||
logger.debug("No running loop, deferring health check task")
|
||||
|
||||
def start_health_checks(self) -> None:
|
||||
"""Start periodic health checks if not already running.
|
||||
|
||||
Call this from the application lifespan when the event loop is available.
|
||||
"""
|
||||
if self._health_task is None and len(self._servers) > 1:
|
||||
loop = asyncio.get_running_loop()
|
||||
self._health_task = loop.create_task(
|
||||
self._periodic_health_check(self._health_check_interval)
|
||||
)
|
||||
logger.info("Started periodic health checks (interval=%ds)", self._health_check_interval)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "llamacpp"
|
||||
|
||||
async def _get_server(self) -> tuple[LlamaCppServer, AsyncOpenAI]:
|
||||
"""Get next healthy server via round-robin."""
|
||||
async with self._lock:
|
||||
healthy = [s for s in self._servers if s.healthy]
|
||||
if not healthy:
|
||||
# All down — try all servers (failover attempt)
|
||||
logger.warning("All llama.cpp servers unhealthy, trying all")
|
||||
healthy = self._servers
|
||||
|
||||
server = healthy[self._current_index % len(healthy)]
|
||||
self._current_index += 1
|
||||
return server, self._clients[server.url]
|
||||
|
||||
def _mark_unhealthy(self, server: LlamaCppServer) -> None:
|
||||
server.error_count += 1
|
||||
if server.healthy:
|
||||
server.healthy = False
|
||||
logger.warning(
|
||||
"llama.cpp server %s marked unhealthy (errors: %d)",
|
||||
server.short_name,
|
||||
server.error_count,
|
||||
)
|
||||
|
||||
def _mark_healthy(self, server: LlamaCppServer) -> None:
|
||||
if not server.healthy:
|
||||
logger.info(
|
||||
"llama.cpp server %s recovered (had %d errors)",
|
||||
server.short_name,
|
||||
server.error_count,
|
||||
)
|
||||
server.error_count = 0
|
||||
server.healthy = True
|
||||
|
||||
async def _periodic_health_check(self, interval: int) -> None:
|
||||
"""Periodically check health of all servers."""
|
||||
while True:
|
||||
await asyncio.sleep(interval)
|
||||
for server in self._servers:
|
||||
client = self._clients[server.url]
|
||||
try:
|
||||
await client.models.list()
|
||||
self._mark_healthy(server)
|
||||
except Exception:
|
||||
self._mark_unhealthy(server)
|
||||
server.last_check = time.time()
|
||||
|
||||
async def _complete_on_server(
|
||||
self,
|
||||
server: LlamaCppServer,
|
||||
client: AsyncOpenAI,
|
||||
messages: list[ChatMessage],
|
||||
model: str,
|
||||
**kwargs: object,
|
||||
) -> CompletionResponse:
|
||||
"""Execute completion on a specific server."""
|
||||
start = time.time()
|
||||
|
||||
response = await client.chat.completions.create(
|
||||
model=model,
|
||||
messages=[m.model_dump(exclude_none=True) for m in messages], # type: ignore[arg-type]
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
elapsed_ms = (time.time() - start) * 1000
|
||||
server.response_times.append(elapsed_ms)
|
||||
if len(server.response_times) > 50:
|
||||
server.response_times = server.response_times[-50:]
|
||||
server.request_count += 1
|
||||
self._mark_healthy(server)
|
||||
|
||||
choices = []
|
||||
for choice in response.choices:
|
||||
choices.append(
|
||||
Choice(
|
||||
index=choice.index,
|
||||
message=ChatMessage(
|
||||
role=choice.message.role,
|
||||
content=choice.message.content,
|
||||
),
|
||||
finish_reason=choice.finish_reason,
|
||||
)
|
||||
)
|
||||
|
||||
usage = None
|
||||
if response.usage:
|
||||
usage = Usage(
|
||||
prompt_tokens=response.usage.prompt_tokens,
|
||||
completion_tokens=response.usage.completion_tokens,
|
||||
total_tokens=response.usage.total_tokens,
|
||||
)
|
||||
|
||||
return CompletionResponse(
|
||||
id=response.id,
|
||||
created=response.created,
|
||||
model=response.model or model,
|
||||
choices=choices,
|
||||
usage=usage,
|
||||
backend=f"{self.name}@{server.short_name}",
|
||||
)
|
||||
|
||||
async def complete(
|
||||
self,
|
||||
messages: list[ChatMessage],
|
||||
model: str,
|
||||
**kwargs: object,
|
||||
) -> CompletionResponse:
|
||||
"""Generate a chat completion with load balancing and failover.
|
||||
|
||||
Tries healthy servers round-robin. On failure, marks server unhealthy
|
||||
and retries on the next available server.
|
||||
|
||||
Args:
|
||||
messages: List of chat messages.
|
||||
model: Model identifier (may be ignored by server).
|
||||
**kwargs: Additional parameters.
|
||||
|
||||
Returns:
|
||||
CompletionResponse: The completion response.
|
||||
|
||||
Raises:
|
||||
CompletionError: If all servers fail.
|
||||
"""
|
||||
tried: set[str] = set()
|
||||
last_error: Exception | None = None
|
||||
|
||||
for _attempt in range(len(self._servers)):
|
||||
server, client = await self._get_server()
|
||||
|
||||
if server.url in tried:
|
||||
continue
|
||||
tried.add(server.url)
|
||||
|
||||
async def _do_complete(
|
||||
_server: LlamaCppServer = server,
|
||||
_client: AsyncOpenAI = client,
|
||||
) -> CompletionResponse:
|
||||
return await self._complete_on_server(
|
||||
_server, _client, messages, model, **kwargs
|
||||
)
|
||||
|
||||
try:
|
||||
return await retry_with_backoff(
|
||||
_do_complete,
|
||||
max_retries=self._settings.max_retries,
|
||||
min_wait=self._settings.retry_min_wait,
|
||||
max_wait=self._settings.retry_max_wait,
|
||||
backend=self.name,
|
||||
model=model,
|
||||
)
|
||||
except (LLMConnectionError, LLMTimeoutError) as e:
|
||||
self._mark_unhealthy(server)
|
||||
last_error = e
|
||||
logger.warning(
|
||||
"llama.cpp server %s failed, trying next: %s",
|
||||
server.short_name,
|
||||
str(e),
|
||||
)
|
||||
continue
|
||||
except Exception:
|
||||
raise
|
||||
|
||||
raise CompletionError(
|
||||
f"All llama.cpp servers failed. Last error: {last_error}",
|
||||
backend=self.name,
|
||||
model=model,
|
||||
)
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
messages: list[ChatMessage],
|
||||
model: str,
|
||||
**kwargs: object,
|
||||
) -> AsyncIterator[CompletionChunk]:
|
||||
"""Generate a streaming chat completion with failover.
|
||||
|
||||
Args:
|
||||
messages: List of chat messages.
|
||||
model: Model identifier.
|
||||
**kwargs: Additional parameters.
|
||||
|
||||
Yields:
|
||||
CompletionChunk: Streaming chunks of the completion.
|
||||
|
||||
Raises:
|
||||
CompletionError: If all servers fail.
|
||||
"""
|
||||
tried: set[str] = set()
|
||||
last_error: Exception | None = None
|
||||
|
||||
for _attempt in range(len(self._servers)):
|
||||
server, client = await self._get_server()
|
||||
|
||||
if server.url in tried:
|
||||
continue
|
||||
tried.add(server.url)
|
||||
|
||||
stream = None
|
||||
try:
|
||||
stream = await client.chat.completions.create(
|
||||
model=model,
|
||||
messages=[m.model_dump(exclude_none=True) for m in messages], # type: ignore[arg-type]
|
||||
stream=True,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
server.request_count += 1
|
||||
self._mark_healthy(server)
|
||||
|
||||
async for chunk in stream:
|
||||
choices = []
|
||||
if chunk.choices:
|
||||
for choice in chunk.choices:
|
||||
delta = Delta(
|
||||
role=getattr(choice.delta, "role", None),
|
||||
content=getattr(choice.delta, "content", None),
|
||||
)
|
||||
choices.append(
|
||||
StreamChoice(
|
||||
index=choice.index,
|
||||
delta=delta,
|
||||
finish_reason=choice.finish_reason,
|
||||
)
|
||||
)
|
||||
|
||||
yield CompletionChunk(
|
||||
id=chunk.id,
|
||||
created=int(time.time()),
|
||||
model=chunk.model or model,
|
||||
choices=choices,
|
||||
)
|
||||
|
||||
# Stream completed successfully
|
||||
return
|
||||
|
||||
except openai.APIConnectionError as e:
|
||||
self._mark_unhealthy(server)
|
||||
last_error = e
|
||||
logger.warning(
|
||||
"llama.cpp server %s stream failed: %s", server.short_name, e
|
||||
)
|
||||
continue
|
||||
except openai.APITimeoutError as e:
|
||||
self._mark_unhealthy(server)
|
||||
last_error = e
|
||||
logger.warning(
|
||||
"llama.cpp server %s stream timeout: %s", server.short_name, e
|
||||
)
|
||||
continue
|
||||
except openai.RateLimitError as e:
|
||||
raise LLMRateLimitError(str(e), backend=self.name) from e
|
||||
except openai.AuthenticationError as e:
|
||||
raise CompletionError(
|
||||
f"Authentication failed: {e}", backend=self.name, model=model
|
||||
) from e
|
||||
except openai.APIStatusError as e:
|
||||
raise CompletionError(str(e), backend=self.name, model=model) from e
|
||||
except Exception as e:
|
||||
raise CompletionError(str(e), backend=self.name, model=model) from e
|
||||
finally:
|
||||
await safe_close_stream(stream, logger)
|
||||
|
||||
raise CompletionError(
|
||||
f"All llama.cpp servers failed for stream. Last error: {last_error}",
|
||||
backend=self.name,
|
||||
model=model,
|
||||
)
|
||||
|
||||
async def list_models(self) -> list[ModelInfo]:
|
||||
"""List models from the first healthy server.
|
||||
|
||||
Returns:
|
||||
list[ModelInfo]: List of available models.
|
||||
|
||||
Raises:
|
||||
ModelListError: If models cannot be retrieved.
|
||||
"""
|
||||
for server in self._servers:
|
||||
if not server.healthy and len(self._servers) > 1:
|
||||
continue
|
||||
client = self._clients[server.url]
|
||||
try:
|
||||
models = await client.models.list()
|
||||
return [
|
||||
ModelInfo(
|
||||
id=m.id,
|
||||
backend=self.name,
|
||||
capabilities=["chat", "vision"],
|
||||
)
|
||||
for m in models.data
|
||||
]
|
||||
except Exception as e:
|
||||
self._mark_unhealthy(server)
|
||||
logger.debug("Failed to list models from %s: %s", server.short_name, e)
|
||||
continue
|
||||
|
||||
raise ModelListError(self.name, "All llama.cpp servers unreachable")
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Check if at least one llama.cpp server is healthy.
|
||||
|
||||
Returns:
|
||||
bool: True if any server is reachable.
|
||||
"""
|
||||
for server in self._servers:
|
||||
client = self._clients[server.url]
|
||||
try:
|
||||
await client.models.list()
|
||||
self._mark_healthy(server)
|
||||
return True
|
||||
except Exception:
|
||||
self._mark_unhealthy(server)
|
||||
continue
|
||||
return False
|
||||
|
||||
def get_servers_status(self) -> list[dict[str, object]]:
|
||||
"""Get status of all servers (for health endpoint)."""
|
||||
return [
|
||||
{
|
||||
"url": s.url,
|
||||
"healthy": s.healthy,
|
||||
"request_count": s.request_count,
|
||||
"error_count": s.error_count,
|
||||
"avg_response_ms": round(s.avg_response_ms, 1),
|
||||
}
|
||||
for s in self._servers
|
||||
]
|
||||
|
|
@ -0,0 +1,113 @@
|
|||
"""Backend registry for managing LLM backend instances."""
|
||||
|
||||
from llm_inference.backends.base import LLMBackend
|
||||
from llm_inference.config import LLMSettings
|
||||
from llm_inference.exceptions import BackendNotAvailableError, BackendNotEnabledError
|
||||
from llm_inference.types import BackendType
|
||||
|
||||
|
||||
class BackendRegistry:
|
||||
"""Registry for managing LLM backend instances.
|
||||
|
||||
The registry initializes and provides access to configured backends.
|
||||
LiteLLM is always available; vLLM and llama.cpp are optional.
|
||||
"""
|
||||
|
||||
def __init__(self, settings: LLMSettings) -> None:
|
||||
"""Initialize the backend registry.
|
||||
|
||||
Args:
|
||||
settings: Application settings.
|
||||
"""
|
||||
self._settings = settings
|
||||
self._backends: dict[BackendType, LLMBackend] = {}
|
||||
self._initialize_backends()
|
||||
|
||||
def _initialize_backends(self) -> None:
|
||||
"""Initialize enabled backends."""
|
||||
# LiteLLM is always available (core dependency)
|
||||
from llm_inference.backends.litellm_backend import LiteLLMBackend
|
||||
|
||||
self._backends[BackendType.LITELLM] = LiteLLMBackend(self._settings)
|
||||
|
||||
# vLLM is optional but MUST work if enabled (fail-fast)
|
||||
if self._settings.enable_vllm:
|
||||
try:
|
||||
from llm_inference.backends.vllm_backend import VLLMBackend
|
||||
|
||||
self._backends[BackendType.VLLM] = VLLMBackend(self._settings)
|
||||
except ImportError as e:
|
||||
raise BackendNotAvailableError(
|
||||
"vllm",
|
||||
f"vLLM backend is enabled but dependencies are not installed. "
|
||||
f"Install with: pip install llm-inference[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 llm_inference.backends.llamacpp_backend import LlamaCppBackend
|
||||
|
||||
self._backends[BackendType.LLAMACPP] = LlamaCppBackend(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 llm-inference[llamacpp]. "
|
||||
f"Original error: {e}",
|
||||
) from e
|
||||
|
||||
def get(self, backend_type: BackendType | str | None = None) -> LLMBackend:
|
||||
"""Get a backend instance.
|
||||
|
||||
Args:
|
||||
backend_type: Backend to retrieve. If None, uses default from settings.
|
||||
|
||||
Returns:
|
||||
LLMBackend: 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,348 @@
|
|||
"""vLLM backend implementation.
|
||||
|
||||
vLLM is a high-throughput and memory-efficient inference engine for LLMs.
|
||||
This backend connects to a vLLM server via its OpenAI-compatible API.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
import httpx
|
||||
|
||||
from llm_inference.backends.base import LLMBackend
|
||||
from llm_inference.config import LLMSettings
|
||||
from llm_inference.exceptions import (
|
||||
CompletionError,
|
||||
LLMConnectionError,
|
||||
LLMRateLimitError,
|
||||
LLMTimeoutError,
|
||||
ModelListError,
|
||||
)
|
||||
from llm_inference.retry import retry_with_backoff
|
||||
from llm_inference.schemas import (
|
||||
CompletionChunk,
|
||||
CompletionResponse,
|
||||
TextChoice,
|
||||
TextCompletionResponse,
|
||||
)
|
||||
from llm_inference.types import (
|
||||
ChatMessage,
|
||||
Choice,
|
||||
Delta,
|
||||
ModelInfo,
|
||||
StreamChoice,
|
||||
Usage,
|
||||
)
|
||||
from llm_inference.utils import safe_close_stream
|
||||
|
||||
try:
|
||||
import openai
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
OPENAI_AVAILABLE = True
|
||||
except ImportError:
|
||||
OPENAI_AVAILABLE = False
|
||||
|
||||
logger = logging.getLogger("llm_inference.backends.vllm")
|
||||
|
||||
|
||||
class VLLMBackend(LLMBackend):
|
||||
"""vLLM backend using OpenAI-compatible API.
|
||||
|
||||
Requires the vLLM server to be running with the OpenAI-compatible
|
||||
API enabled.
|
||||
|
||||
Example:
|
||||
Start vLLM server:
|
||||
```bash
|
||||
python -m vllm.entrypoints.openai.api_server \
|
||||
--model meta-llama/Llama-2-7b-chat-hf \
|
||||
--host 0.0.0.0 --port 8101
|
||||
```
|
||||
|
||||
Configure:
|
||||
```bash
|
||||
export LLM_ENABLE_VLLM=true
|
||||
export LLM_VLLM_BASE_URL=http://localhost:8101
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, settings: LLMSettings) -> 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 llm-inference[vllm]"
|
||||
)
|
||||
|
||||
self._settings = settings
|
||||
|
||||
# Configure httpx timeout
|
||||
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, # We handle retries ourselves
|
||||
)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
"""Backend identifier name."""
|
||||
return "vllm"
|
||||
|
||||
async def complete(
|
||||
self,
|
||||
messages: list[ChatMessage],
|
||||
model: str,
|
||||
**kwargs: object,
|
||||
) -> CompletionResponse:
|
||||
"""Generate a chat completion via vLLM.
|
||||
|
||||
Args:
|
||||
messages: List of chat messages.
|
||||
model: Model identifier.
|
||||
**kwargs: Additional parameters.
|
||||
|
||||
Returns:
|
||||
CompletionResponse: The completion response.
|
||||
|
||||
Raises:
|
||||
CompletionError: If the request fails after retries.
|
||||
LLMRateLimitError: If rate limited.
|
||||
LLMTimeoutError: If the request times out.
|
||||
"""
|
||||
|
||||
async def _do_complete() -> CompletionResponse:
|
||||
response = await self._client.chat.completions.create(
|
||||
model=model,
|
||||
messages=[m.model_dump(exclude_none=True) for m in messages], # type: ignore[arg-type]
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
choices = []
|
||||
for choice in response.choices:
|
||||
choices.append(
|
||||
Choice(
|
||||
index=choice.index,
|
||||
message=ChatMessage(
|
||||
role=choice.message.role,
|
||||
content=choice.message.content,
|
||||
),
|
||||
finish_reason=choice.finish_reason,
|
||||
)
|
||||
)
|
||||
|
||||
usage = None
|
||||
if response.usage:
|
||||
usage = Usage(
|
||||
prompt_tokens=response.usage.prompt_tokens,
|
||||
completion_tokens=response.usage.completion_tokens,
|
||||
total_tokens=response.usage.total_tokens,
|
||||
)
|
||||
|
||||
return CompletionResponse(
|
||||
id=response.id,
|
||||
created=response.created,
|
||||
model=response.model,
|
||||
choices=choices,
|
||||
usage=usage,
|
||||
backend=self.name,
|
||||
)
|
||||
|
||||
return await retry_with_backoff(
|
||||
_do_complete,
|
||||
max_retries=self._settings.max_retries,
|
||||
min_wait=self._settings.retry_min_wait,
|
||||
max_wait=self._settings.retry_max_wait,
|
||||
backend=self.name,
|
||||
model=model,
|
||||
)
|
||||
|
||||
async def text_complete(
|
||||
self,
|
||||
prompt: str | list[str],
|
||||
model: str,
|
||||
**kwargs: object,
|
||||
) -> TextCompletionResponse:
|
||||
"""Generate a legacy text completion via vLLM (/v1/completions)."""
|
||||
|
||||
async def _do_text() -> TextCompletionResponse:
|
||||
response = await self._client.completions.create(
|
||||
model=model,
|
||||
prompt=prompt, # type: ignore[arg-type]
|
||||
**kwargs,
|
||||
)
|
||||
choices = [
|
||||
TextChoice(
|
||||
index=c.index, text=c.text, finish_reason=c.finish_reason
|
||||
)
|
||||
for c in response.choices
|
||||
]
|
||||
usage = None
|
||||
if response.usage:
|
||||
usage = Usage(
|
||||
prompt_tokens=response.usage.prompt_tokens,
|
||||
completion_tokens=response.usage.completion_tokens,
|
||||
total_tokens=response.usage.total_tokens,
|
||||
)
|
||||
return TextCompletionResponse(
|
||||
id=response.id,
|
||||
created=response.created,
|
||||
model=response.model,
|
||||
choices=choices,
|
||||
usage=usage,
|
||||
backend=self.name,
|
||||
)
|
||||
|
||||
return await retry_with_backoff(
|
||||
_do_text,
|
||||
max_retries=self._settings.max_retries,
|
||||
min_wait=self._settings.retry_min_wait,
|
||||
max_wait=self._settings.retry_max_wait,
|
||||
backend=self.name,
|
||||
model=model,
|
||||
)
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
messages: list[ChatMessage],
|
||||
model: str,
|
||||
**kwargs: object,
|
||||
) -> AsyncIterator[CompletionChunk]:
|
||||
"""Generate a streaming chat completion via vLLM.
|
||||
|
||||
Args:
|
||||
messages: List of chat messages.
|
||||
model: Model identifier.
|
||||
**kwargs: Additional parameters.
|
||||
|
||||
Yields:
|
||||
CompletionChunk: Streaming chunks of the completion.
|
||||
|
||||
Raises:
|
||||
CompletionError: If the request fails.
|
||||
LLMRateLimitError: If rate limited.
|
||||
LLMTimeoutError: If the request times out.
|
||||
"""
|
||||
stream = None
|
||||
try:
|
||||
stream = await self._client.chat.completions.create(
|
||||
model=model,
|
||||
messages=[m.model_dump(exclude_none=True) for m in messages], # type: ignore[arg-type]
|
||||
stream=True,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
async for chunk in stream:
|
||||
choices = []
|
||||
if chunk.choices:
|
||||
for choice in chunk.choices:
|
||||
delta = Delta(
|
||||
role=getattr(choice.delta, "role", None),
|
||||
content=getattr(choice.delta, "content", None),
|
||||
)
|
||||
choices.append(
|
||||
StreamChoice(
|
||||
index=choice.index,
|
||||
delta=delta,
|
||||
finish_reason=choice.finish_reason,
|
||||
)
|
||||
)
|
||||
|
||||
yield CompletionChunk(
|
||||
id=chunk.id,
|
||||
created=int(time.time()),
|
||||
model=chunk.model or model,
|
||||
choices=choices,
|
||||
)
|
||||
|
||||
except openai.RateLimitError as e:
|
||||
logger.warning("Rate limited during stream: %s", str(e))
|
||||
raise LLMRateLimitError(str(e), backend=self.name) from e
|
||||
except openai.APITimeoutError as e:
|
||||
logger.warning("Timeout during stream: %s", str(e))
|
||||
raise LLMTimeoutError(str(e), backend=self.name) from e
|
||||
except openai.APIConnectionError as e:
|
||||
logger.error("Connection error during stream: %s", str(e))
|
||||
raise LLMConnectionError(self.name, reason=str(e)) from e
|
||||
except openai.AuthenticationError as e:
|
||||
logger.error("Authentication failed: %s", str(e))
|
||||
raise CompletionError(
|
||||
f"Authentication failed: {e}", backend=self.name, model=model
|
||||
) from e
|
||||
except openai.APIStatusError as e:
|
||||
logger.error("API error during stream: %s", str(e))
|
||||
raise CompletionError(str(e), backend=self.name, model=model) from e
|
||||
except Exception as e:
|
||||
logger.error("Unexpected error during stream: %s", str(e))
|
||||
raise CompletionError(str(e), backend=self.name, model=model) from e
|
||||
finally:
|
||||
# Ensure stream is properly closed on cancellation or error
|
||||
await safe_close_stream(stream, logger)
|
||||
|
||||
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,
|
||||
capabilities=["chat"],
|
||||
)
|
||||
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
|
||||
140
ai_platform/modules/llm-inference/src/llm_inference/cli.py
Normal file
140
ai_platform/modules/llm-inference/src/llm_inference/cli.py
Normal file
|
|
@ -0,0 +1,140 @@
|
|||
"""CLI entry point for the LLM inference server."""
|
||||
|
||||
import argparse
|
||||
import signal
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
import uvicorn
|
||||
|
||||
from llm_inference.config import SettingsCache
|
||||
from llm_inference.logging import configure_logging, get_logger
|
||||
|
||||
|
||||
class GracefulShutdown:
|
||||
"""Handles graceful shutdown on signals.
|
||||
|
||||
Allows in-flight requests to complete before exiting.
|
||||
"""
|
||||
|
||||
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.
|
||||
|
||||
Args:
|
||||
signum: Signal number.
|
||||
frame: Current stack frame (unused).
|
||||
"""
|
||||
if self._shutdown_requested:
|
||||
# Second signal - force exit
|
||||
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.
|
||||
|
||||
Args:
|
||||
server: Uvicorn server instance.
|
||||
"""
|
||||
self._server = server
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Run the LLM inference server with graceful shutdown support."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="LLM Inference API Server",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--host",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Host to bind to (default: from LLM_HOST env or 0.0.0.0)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--port",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Port to bind to (default: from LLM_PORT env or 14011)",
|
||||
)
|
||||
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()
|
||||
|
||||
# Get settings for defaults
|
||||
settings = SettingsCache.get()
|
||||
|
||||
# Configure logging early
|
||||
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 LLM Inference API Server on %s:%d", host, port)
|
||||
|
||||
if args.reload:
|
||||
# Development mode - use simple uvicorn.run
|
||||
logger.info("Development mode - auto-reload enabled")
|
||||
uvicorn.run(
|
||||
"llm_inference.api.app:create_app",
|
||||
factory=True,
|
||||
host=host,
|
||||
port=port,
|
||||
workers=1,
|
||||
reload=True,
|
||||
)
|
||||
else:
|
||||
# Production mode - use server with graceful shutdown
|
||||
shutdown_handler = GracefulShutdown()
|
||||
shutdown_handler.register_signals()
|
||||
|
||||
config = uvicorn.Config(
|
||||
"llm_inference.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()
|
||||
357
ai_platform/modules/llm-inference/src/llm_inference/client.py
Normal file
357
ai_platform/modules/llm-inference/src/llm_inference/client.py
Normal file
|
|
@ -0,0 +1,357 @@
|
|||
"""High-level LLM client for unified inference."""
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from time import perf_counter
|
||||
|
||||
from llm_inference import metrics
|
||||
from llm_inference.backends import BackendRegistry
|
||||
from llm_inference.config import LLMSettings, SettingsCache
|
||||
from llm_inference.exceptions import (
|
||||
BackendNotAvailableError,
|
||||
BackendNotEnabledError,
|
||||
CompletionError,
|
||||
LLMRateLimitError,
|
||||
LLMTimeoutError,
|
||||
)
|
||||
from llm_inference.image_processing import process_messages
|
||||
from llm_inference.schemas import (
|
||||
CompletionChunk,
|
||||
CompletionResponse,
|
||||
TextCompletionResponse,
|
||||
)
|
||||
from llm_inference.types import BackendType, ChatMessage, ModelInfo
|
||||
|
||||
|
||||
class LLMClient:
|
||||
"""High-level client for LLM 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 llm_inference import LLMClient
|
||||
|
||||
client = LLMClient()
|
||||
|
||||
# Simple completion
|
||||
response = await client.complete(
|
||||
messages=[{"role": "user", "content": "Hello!"}],
|
||||
model="gpt-3.5-turbo",
|
||||
)
|
||||
print(response.choices[0].message.content)
|
||||
|
||||
# Streaming
|
||||
async for chunk in client.stream(
|
||||
messages=[{"role": "user", "content": "Tell me a story"}],
|
||||
model="gpt-4",
|
||||
):
|
||||
if chunk.choices[0].delta.content:
|
||||
print(chunk.choices[0].delta.content, end="")
|
||||
|
||||
# Override backend per-request
|
||||
response = await client.complete(
|
||||
messages=[{"role": "user", "content": "Hello!"}],
|
||||
model="meta-llama/Llama-2-7b-chat-hf",
|
||||
backend=BackendType.VLLM,
|
||||
)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, settings: LLMSettings | None = None) -> None:
|
||||
"""Initialize the LLM 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 complete(
|
||||
self,
|
||||
messages: list[dict[str, str] | ChatMessage],
|
||||
model: str | None = None,
|
||||
backend: BackendType | str | None = None,
|
||||
**kwargs: object,
|
||||
) -> CompletionResponse:
|
||||
"""Generate a chat completion.
|
||||
|
||||
Args:
|
||||
messages: List of chat messages. Can be dicts or ChatMessage objects.
|
||||
model: Model identifier. Uses default if not specified.
|
||||
backend: Backend to use. If not specified, auto-detects from model
|
||||
or uses default.
|
||||
**kwargs: Additional parameters (temperature, max_tokens, etc.).
|
||||
|
||||
Returns:
|
||||
CompletionResponse: The completion response.
|
||||
|
||||
Raises:
|
||||
BackendNotAvailableError: If the requested backend is not available.
|
||||
CompletionError: If the completion request fails.
|
||||
"""
|
||||
parsed_messages = self._parse_messages(messages)
|
||||
parsed_messages = await process_messages(parsed_messages)
|
||||
model = self._resolve_model(model or self._settings.default_model)
|
||||
explicit = backend is not None
|
||||
if backend is None:
|
||||
backend = await self._resolve_backend_for_model(model)
|
||||
|
||||
cascade = self._build_cascade(backend, explicit)
|
||||
last_error: Exception | None = None
|
||||
for backend_type in cascade:
|
||||
name = backend_type.value
|
||||
start = perf_counter()
|
||||
metrics.LLM_INFLIGHT.labels(backend=name).inc()
|
||||
try:
|
||||
response = await self._registry.get(backend_type).complete(
|
||||
parsed_messages, model, **kwargs
|
||||
)
|
||||
metrics.observe_success(
|
||||
model, response.backend or name, perf_counter() - start,
|
||||
response.usage,
|
||||
)
|
||||
return response
|
||||
except (
|
||||
BackendNotAvailableError,
|
||||
BackendNotEnabledError,
|
||||
CompletionError,
|
||||
LLMTimeoutError,
|
||||
LLMRateLimitError,
|
||||
) as exc:
|
||||
metrics.observe_error(model, name)
|
||||
last_error = exc
|
||||
continue
|
||||
finally:
|
||||
metrics.LLM_INFLIGHT.labels(backend=name).dec()
|
||||
# All backends in the cascade failed — surface the last error.
|
||||
raise last_error if last_error else CompletionError("no backend available")
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
messages: list[dict[str, str] | ChatMessage],
|
||||
model: str | None = None,
|
||||
backend: BackendType | str | None = None,
|
||||
**kwargs: object,
|
||||
) -> AsyncIterator[CompletionChunk]:
|
||||
"""Generate a streaming chat completion.
|
||||
|
||||
Args:
|
||||
messages: List of chat messages.
|
||||
model: Model identifier. Uses default if not specified.
|
||||
backend: Backend to use. Uses default if not specified.
|
||||
**kwargs: Additional parameters.
|
||||
|
||||
Yields:
|
||||
CompletionChunk: Streaming chunks of the completion.
|
||||
|
||||
Raises:
|
||||
BackendNotAvailableError: If the requested backend is not available.
|
||||
CompletionError: If the completion request fails.
|
||||
"""
|
||||
parsed_messages = self._parse_messages(messages)
|
||||
parsed_messages = await process_messages(parsed_messages)
|
||||
model = self._resolve_model(model or self._settings.default_model)
|
||||
if backend is None:
|
||||
backend = await self._resolve_backend_for_model(model)
|
||||
backend_instance = self._registry.get(backend)
|
||||
|
||||
async for chunk in backend_instance.stream(parsed_messages, model, **kwargs):
|
||||
yield chunk
|
||||
|
||||
async def text_complete(
|
||||
self,
|
||||
prompt: str | list[str],
|
||||
model: str | None = None,
|
||||
backend: BackendType | str | None = None,
|
||||
**kwargs: object,
|
||||
) -> TextCompletionResponse:
|
||||
"""Generate a legacy text completion (/v1/completions).
|
||||
|
||||
Resolves the model alias and target backend, then delegates to the
|
||||
backend's ``text_complete``. Backends that do not support text
|
||||
completion raise NotImplementedError.
|
||||
"""
|
||||
model = self._resolve_model(model or self._settings.default_model)
|
||||
if backend is None:
|
||||
backend = await self._resolve_backend_for_model(model)
|
||||
backend_instance = self._registry.get(backend)
|
||||
return await backend_instance.text_complete(prompt, model, **kwargs)
|
||||
|
||||
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
|
||||
|
||||
async def load_model(
|
||||
self,
|
||||
model: str,
|
||||
backend: BackendType | str,
|
||||
) -> bool:
|
||||
"""Load a model on the specified backend.
|
||||
|
||||
This is only supported by local backends (vLLM, llama.cpp).
|
||||
|
||||
Args:
|
||||
model: Model identifier to load.
|
||||
backend: Backend to load the model on.
|
||||
|
||||
Returns:
|
||||
bool: True if the model was loaded successfully.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If the backend doesn't support model loading.
|
||||
ModelLoadError: If loading fails.
|
||||
"""
|
||||
backend_instance = self._registry.get(backend)
|
||||
return await backend_instance.load_model(model)
|
||||
|
||||
async def unload_model(
|
||||
self,
|
||||
model: str,
|
||||
backend: BackendType | str,
|
||||
) -> bool:
|
||||
"""Unload a model from the specified backend.
|
||||
|
||||
This is only supported by local backends (vLLM, llama.cpp).
|
||||
|
||||
Args:
|
||||
model: Model identifier to unload.
|
||||
backend: Backend to unload the model from.
|
||||
|
||||
Returns:
|
||||
bool: True if the model was unloaded successfully.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If the backend doesn't support model unloading.
|
||||
"""
|
||||
backend_instance = self._registry.get(backend)
|
||||
return await backend_instance.unload_model(model)
|
||||
|
||||
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
|
||||
|
||||
def _resolve_model(self, model: str) -> str:
|
||||
"""Resolve a friendly model alias to its real served id.
|
||||
|
||||
Returns the configured alias target, or the input unchanged when no
|
||||
alias is defined (passthrough).
|
||||
"""
|
||||
return self._settings.model_aliases.get(model, model)
|
||||
|
||||
def _build_cascade(
|
||||
self, primary: BackendType | str | None, explicit: bool
|
||||
) -> list[BackendType]:
|
||||
"""Build the ordered backend cascade for a request.
|
||||
|
||||
Returns ``[primary]`` when the caller pinned a backend or fallback is
|
||||
disabled; otherwise ``[primary, *other enabled backends]`` in
|
||||
``fallback_order`` (offer §1.6).
|
||||
"""
|
||||
if primary is None:
|
||||
primary = BackendType(self._settings.default_backend)
|
||||
elif isinstance(primary, str):
|
||||
primary = BackendType(primary)
|
||||
|
||||
if explicit or not self._settings.enable_fallback:
|
||||
return [primary]
|
||||
|
||||
enabled = set(self._registry.list_backends())
|
||||
chain: list[BackendType] = [primary] if primary in enabled else []
|
||||
for name in self._settings.fallback_order:
|
||||
try:
|
||||
bt = BackendType(name)
|
||||
except ValueError:
|
||||
continue
|
||||
if bt in enabled and bt != primary:
|
||||
chain.append(bt)
|
||||
return chain or [primary]
|
||||
|
||||
async def _resolve_backend_for_model(self, model: str) -> BackendType | None:
|
||||
"""Auto-detect the correct backend for a model.
|
||||
|
||||
Queries each enabled local backend's model list to find which one
|
||||
serves the requested model. Returns None to use default.
|
||||
|
||||
Args:
|
||||
model: Model identifier to look up.
|
||||
|
||||
Returns:
|
||||
BackendType if found, None for default.
|
||||
"""
|
||||
for backend_type in self._registry.list_backends():
|
||||
if backend_type == BackendType.LITELLM:
|
||||
continue
|
||||
try:
|
||||
models = await self._registry.get(backend_type).list_models()
|
||||
if any(m.id == model for m in models):
|
||||
return backend_type
|
||||
except Exception:
|
||||
continue
|
||||
return None
|
||||
|
||||
def _parse_messages(
|
||||
self,
|
||||
messages: list[dict[str, str] | ChatMessage],
|
||||
) -> list[ChatMessage]:
|
||||
"""Parse messages into ChatMessage objects.
|
||||
|
||||
Args:
|
||||
messages: List of messages as dicts or ChatMessage objects.
|
||||
|
||||
Returns:
|
||||
list[ChatMessage]: List of ChatMessage objects.
|
||||
"""
|
||||
parsed: list[ChatMessage] = []
|
||||
for msg in messages:
|
||||
if isinstance(msg, ChatMessage):
|
||||
parsed.append(msg)
|
||||
else:
|
||||
parsed.append(ChatMessage.model_validate(msg))
|
||||
return parsed
|
||||
337
ai_platform/modules/llm-inference/src/llm_inference/config.py
Normal file
337
ai_platform/modules/llm-inference/src/llm_inference/config.py
Normal file
|
|
@ -0,0 +1,337 @@
|
|||
"""Configuration management using Pydantic V2 Settings."""
|
||||
|
||||
import threading
|
||||
import warnings
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import Field, field_validator, model_validator
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
|
||||
class LLMSettings(BaseSettings):
|
||||
"""LLM Inference module configuration.
|
||||
|
||||
All settings can be configured via environment variables with the LLM_ prefix.
|
||||
|
||||
Required environment variables (no defaults - fail-fast):
|
||||
LLM_DEFAULT_BACKEND: Backend to use (litellm, vllm, llamacpp)
|
||||
LLM_ENABLE_VLLM: Enable vLLM backend (true/false)
|
||||
LLM_ENABLE_LLAMACPP: Enable llama.cpp backend (true/false)
|
||||
|
||||
Example:
|
||||
LLM_DEFAULT_BACKEND=litellm
|
||||
LLM_ENABLE_VLLM=false
|
||||
LLM_ENABLE_LLAMACPP=false
|
||||
LLM_PORT=14011
|
||||
"""
|
||||
|
||||
model_config = SettingsConfigDict(
|
||||
env_prefix="LLM_",
|
||||
env_file=".env",
|
||||
env_file_encoding="utf-8",
|
||||
extra="forbid", # Reject unknown fields to catch typos
|
||||
)
|
||||
|
||||
# ==========================================================================
|
||||
# REQUIRED fields (no defaults - fail-fast)
|
||||
# ==========================================================================
|
||||
|
||||
default_backend: Literal["litellm", "vllm", "llamacpp"] = Field(
|
||||
description="Default backend to use for inference (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)",
|
||||
)
|
||||
|
||||
# ==========================================================================
|
||||
# Optional fields with sensible defaults
|
||||
# ==========================================================================
|
||||
|
||||
# Default model selection — the platform's primary model (offer: Qwen3.5).
|
||||
default_model: str = Field(
|
||||
default="qwen3.5",
|
||||
description="Default model alias used when a request omits `model`",
|
||||
)
|
||||
# Friendly alias → real served model id. Lets callers use a stable name
|
||||
# (e.g. "qwen3.5") regardless of the backend's served model id. Set via
|
||||
# LLM_MODEL_ALIASES='{"qwen3.5":"Qwen/Qwen3.5-35B-A3B"}'. Empty = passthrough
|
||||
# (works when vLLM is started with --served-model-name qwen3.5).
|
||||
model_aliases: dict[str, str] = Field(
|
||||
default_factory=dict,
|
||||
description="Map of friendly model alias to the real served model id",
|
||||
)
|
||||
# Cross-backend fallback cascade (offer §1.6): on a backend failure, retry
|
||||
# the request on the next enabled backend in this order. Disabled when the
|
||||
# caller pins an explicit backend.
|
||||
enable_fallback: bool = Field(
|
||||
default=True,
|
||||
description="Enable cross-backend fallback cascade on backend failure",
|
||||
)
|
||||
fallback_order: list[str] = Field(
|
||||
default_factory=lambda: ["vllm", "llamacpp", "litellm"],
|
||||
description="Backend order tried in the fallback cascade",
|
||||
)
|
||||
|
||||
# API server settings
|
||||
host: str = Field(
|
||||
default="0.0.0.0",
|
||||
description="Host to bind the server to",
|
||||
)
|
||||
port: int = Field(
|
||||
default=14011,
|
||||
description="Port to bind the server to (Prod LLM API: 14011)",
|
||||
)
|
||||
external_url: str = Field(
|
||||
description="External URL for OpenAPI spec (e.g., http://10.11.10.42:14011). REQUIRED.",
|
||||
)
|
||||
|
||||
# LiteLLM / Provider API keys (optional but at least one needed for litellm)
|
||||
openrouter_api_key: str | None = Field(
|
||||
default=None,
|
||||
description="OpenRouter API key",
|
||||
)
|
||||
openai_api_key: str | None = Field(
|
||||
default=None,
|
||||
description="OpenAI API key",
|
||||
)
|
||||
anthropic_api_key: str | None = Field(
|
||||
default=None,
|
||||
description="Anthropic API key",
|
||||
)
|
||||
|
||||
# vLLM settings
|
||||
vllm_base_url: str = Field(
|
||||
default="http://localhost:14001",
|
||||
description="Base URL for vLLM 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:8080",
|
||||
description="Base URL for llama.cpp server (single server, legacy)",
|
||||
)
|
||||
llamacpp_base_urls: list[str] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Comma-separated list of llama.cpp server URLs for load balancing. "
|
||||
"If set, overrides llamacpp_base_url. "
|
||||
"Example: http://10.11.10.18:14001"
|
||||
),
|
||||
)
|
||||
llamacpp_health_check_interval: int = Field(
|
||||
default=30,
|
||||
ge=5,
|
||||
description="Health check interval in seconds for llama.cpp backends",
|
||||
)
|
||||
|
||||
@field_validator("llamacpp_base_urls", mode="before")
|
||||
@classmethod
|
||||
def parse_llamacpp_base_urls(cls, v: str | list[str] | None) -> list[str] | None:
|
||||
"""Parse comma-separated URLs into a list."""
|
||||
if v is None or v == "":
|
||||
return None
|
||||
if isinstance(v, str):
|
||||
urls = [u.strip().rstrip("/") for u in v.split(",") if u.strip()]
|
||||
return urls if urls else None
|
||||
return [u.strip().rstrip("/") for u in v if u.strip()]
|
||||
|
||||
@property
|
||||
def llamacpp_urls(self) -> list[str]:
|
||||
"""Get all llama.cpp server URLs (from list or single)."""
|
||||
if self.llamacpp_base_urls:
|
||||
return self.llamacpp_base_urls
|
||||
return [self.llamacpp_base_url.rstrip("/")]
|
||||
|
||||
# Model storage
|
||||
models_dir: str = Field(
|
||||
default="/models",
|
||||
description="Directory for local model storage",
|
||||
)
|
||||
|
||||
# ==========================================================================
|
||||
# Timeout and retry settings
|
||||
# ==========================================================================
|
||||
|
||||
request_timeout: float = Field(
|
||||
default=120.0,
|
||||
ge=1.0,
|
||||
description="Default timeout in seconds for completion requests",
|
||||
)
|
||||
connect_timeout: float = Field(
|
||||
default=10.0,
|
||||
ge=1.0,
|
||||
description="Timeout in seconds for establishing connections",
|
||||
)
|
||||
max_retries: int = Field(
|
||||
default=3,
|
||||
ge=0,
|
||||
description="Maximum number of retries for transient failures",
|
||||
)
|
||||
retry_min_wait: float = Field(
|
||||
default=1.0,
|
||||
ge=0.1,
|
||||
description="Minimum wait time in seconds between retries",
|
||||
)
|
||||
retry_max_wait: float = Field(
|
||||
default=60.0,
|
||||
ge=1.0,
|
||||
description="Maximum wait time in seconds between retries",
|
||||
)
|
||||
|
||||
# ==========================================================================
|
||||
# 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=10.0,
|
||||
ge=0.1,
|
||||
description="Rate limit: requests per second",
|
||||
)
|
||||
rate_limit_burst: int = Field(
|
||||
default=20,
|
||||
ge=1,
|
||||
description="Rate limit: burst capacity",
|
||||
)
|
||||
max_concurrent_completions: int = Field(
|
||||
default=10,
|
||||
ge=1,
|
||||
description="Maximum concurrent completion 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 (e.g., http://didiAI-dashboard:51300). "
|
||||
"When set, runtime_config polls /api/config every 30s for live "
|
||||
"overrides. Unset disables polling — module uses env-only config."
|
||||
),
|
||||
)
|
||||
|
||||
@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)
|
||||
|
||||
# ==========================================================================
|
||||
# Validators
|
||||
# ==========================================================================
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_api_keys(self) -> "LLMSettings":
|
||||
"""Warn if no API keys are set when using litellm backend."""
|
||||
if self.default_backend == "litellm":
|
||||
has_any_key = any(
|
||||
[
|
||||
self.openrouter_api_key,
|
||||
self.openai_api_key,
|
||||
self.anthropic_api_key,
|
||||
]
|
||||
)
|
||||
if not has_any_key:
|
||||
warnings.warn(
|
||||
"No API keys configured for litellm backend. "
|
||||
"Set at least one of: OPENROUTER_API_KEY, "
|
||||
"OPENAI_API_KEY, ANTHROPIC_API_KEY",
|
||||
UserWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
class SettingsCache:
|
||||
"""Thread-safe settings cache that can be cleared for testing."""
|
||||
|
||||
_instance: LLMSettings | None = None
|
||||
_lock: threading.Lock = threading.Lock()
|
||||
|
||||
@classmethod
|
||||
def get(cls) -> LLMSettings:
|
||||
"""Get or create the settings instance.
|
||||
|
||||
Thread-safe singleton pattern. Always acquires lock for correctness.
|
||||
Performance impact is negligible for singleton access patterns.
|
||||
|
||||
Returns:
|
||||
LLMSettings: Application settings loaded from environment.
|
||||
"""
|
||||
with cls._lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = LLMSettings()
|
||||
return cls._instance
|
||||
|
||||
@classmethod
|
||||
def clear(cls) -> None:
|
||||
"""Clear the cached settings instance.
|
||||
|
||||
Use this in tests to reset settings between test cases.
|
||||
"""
|
||||
with cls._lock:
|
||||
cls._instance = None
|
||||
|
||||
@classmethod
|
||||
def set(cls, settings: LLMSettings) -> None:
|
||||
"""Set a specific settings instance.
|
||||
|
||||
Use this in tests to inject mock settings.
|
||||
|
||||
Args:
|
||||
settings: Settings instance to use.
|
||||
"""
|
||||
with cls._lock:
|
||||
cls._instance = settings
|
||||
|
||||
|
||||
def get_settings() -> LLMSettings:
|
||||
"""Get cached settings instance.
|
||||
|
||||
Returns:
|
||||
LLMSettings: Application settings loaded from environment.
|
||||
"""
|
||||
return SettingsCache.get()
|
||||
|
|
@ -0,0 +1,137 @@
|
|||
"""Custom exceptions for LLM inference."""
|
||||
|
||||
|
||||
class LLMInferenceError(Exception):
|
||||
"""Base exception for LLM inference errors."""
|
||||
|
||||
|
||||
class AuthenticationError(LLMInferenceError):
|
||||
"""Raised when authentication fails."""
|
||||
|
||||
def __init__(self, message: str = "Authentication required") -> None:
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class BackendNotAvailableError(LLMInferenceError):
|
||||
"""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(LLMInferenceError):
|
||||
"""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 LLM_ENABLE_{backend.upper()}=true to enable it."
|
||||
)
|
||||
|
||||
|
||||
class ModelNotFoundError(LLMInferenceError):
|
||||
"""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 CompletionError(LLMInferenceError):
|
||||
"""Raised when a completion 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 ModelLoadError(LLMInferenceError):
|
||||
"""Raised when loading a model fails."""
|
||||
|
||||
def __init__(self, model: str, backend: str, reason: str | None = None) -> None:
|
||||
self.model = model
|
||||
self.backend = backend
|
||||
self.reason = reason
|
||||
message = f"Failed to load model '{model}' on backend '{backend}'"
|
||||
if reason:
|
||||
message += f": {reason}"
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class ModelListError(LLMInferenceError):
|
||||
"""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 LLMRateLimitError(LLMInferenceError):
|
||||
"""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 LLMTimeoutError(LLMInferenceError):
|
||||
"""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 LLMConnectionError(LLMInferenceError):
|
||||
"""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)
|
||||
|
|
@ -0,0 +1,185 @@
|
|||
"""Image processing for multimodal requests.
|
||||
|
||||
Downloads image URLs and converts them to base64 data URIs so that all
|
||||
backends (vLLM, llama.cpp, LiteLLM) receive a uniform format.
|
||||
|
||||
Supports:
|
||||
1. OpenAI format: {"type": "image_url", "image_url": {"url": "https://..."}}
|
||||
2. Already base64: {"type": "image_url", "image_url": {"url": "data:image/...;base64,..."}}
|
||||
→ passed through unchanged.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import logging
|
||||
import mimetypes
|
||||
import re
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
|
||||
from llm_inference.types import ChatMessage
|
||||
|
||||
logger = logging.getLogger("llm_inference.image_processing")
|
||||
|
||||
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"}
|
||||
IMAGE_URL_PATTERN = re.compile(
|
||||
r"https?://[^\s<>\"]+\.(?:jpg|jpeg|png|gif|webp|bmp)", re.IGNORECASE
|
||||
)
|
||||
|
||||
# Shared client for image downloads (reused across requests)
|
||||
_http_client: httpx.AsyncClient | None = None
|
||||
|
||||
|
||||
def _get_http_client() -> httpx.AsyncClient:
|
||||
global _http_client
|
||||
if _http_client is None:
|
||||
_http_client = httpx.AsyncClient(
|
||||
timeout=30,
|
||||
follow_redirects=True,
|
||||
headers={"User-Agent": "LLMInference/1.0"},
|
||||
)
|
||||
return _http_client
|
||||
|
||||
|
||||
async def close_http_client() -> None:
|
||||
"""Close the shared HTTP client. Call on application shutdown."""
|
||||
global _http_client
|
||||
if _http_client is not None:
|
||||
await _http_client.aclose()
|
||||
_http_client = None
|
||||
|
||||
|
||||
def _guess_mime_type(url: str, content_type: str = "") -> str:
|
||||
"""Guess MIME type from URL or Content-Type header."""
|
||||
if "jpeg" in content_type or "jpg" in content_type:
|
||||
return "image/jpeg"
|
||||
if "png" in content_type:
|
||||
return "image/png"
|
||||
if "gif" in content_type:
|
||||
return "image/gif"
|
||||
if "webp" in content_type:
|
||||
return "image/webp"
|
||||
if "bmp" in content_type:
|
||||
return "image/bmp"
|
||||
|
||||
ext = urlparse(url).path.rsplit(".", 1)[-1].lower() if "." in url else ""
|
||||
mime = mimetypes.guess_type(f"file.{ext}")[0]
|
||||
return mime or "image/jpeg"
|
||||
|
||||
|
||||
async def _download_and_encode(url: str) -> str:
|
||||
"""Download image URL and return as base64 data URI."""
|
||||
client = _get_http_client()
|
||||
resp = await client.get(url)
|
||||
resp.raise_for_status()
|
||||
|
||||
content_type = resp.headers.get("content-type", "")
|
||||
mime_type = _guess_mime_type(url, content_type)
|
||||
b64 = base64.b64encode(resp.content).decode("utf-8")
|
||||
|
||||
logger.info(
|
||||
"Downloaded image: %s (%d bytes, %s)", url, len(resp.content), mime_type
|
||||
)
|
||||
return f"data:{mime_type};base64,{b64}"
|
||||
|
||||
|
||||
async def _process_content_item(item: dict[str, object]) -> dict[str, object]:
|
||||
"""Process a single content item, downloading image URLs if needed."""
|
||||
if not isinstance(item, dict):
|
||||
return item
|
||||
|
||||
if item.get("type") != "image_url":
|
||||
return item
|
||||
|
||||
image_url = item.get("image_url", {})
|
||||
url = image_url.get("url", "") if isinstance(image_url, dict) else str(image_url)
|
||||
|
||||
# Already base64 — pass through
|
||||
if url.startswith("data:"):
|
||||
return item
|
||||
|
||||
# Not an HTTP URL — pass through
|
||||
if not url.startswith(("http://", "https://")):
|
||||
return item
|
||||
|
||||
try:
|
||||
data_uri = await _download_and_encode(url)
|
||||
return {
|
||||
"type": "image_url",
|
||||
"image_url": {"url": data_uri},
|
||||
}
|
||||
except Exception as e:
|
||||
logger.warning("Failed to download image %s: %s", url, e)
|
||||
return item # Return original on failure
|
||||
|
||||
|
||||
async def _process_text_with_urls(text: str) -> str | list[dict[str, object]]:
|
||||
"""Detect image URLs in plain text and convert to multimodal format."""
|
||||
urls = IMAGE_URL_PATTERN.findall(text)
|
||||
if not urls:
|
||||
return text
|
||||
|
||||
# Build multimodal content array
|
||||
content: list[dict[str, object]] = []
|
||||
|
||||
# Strip URLs from text
|
||||
clean_text = text
|
||||
for url in urls:
|
||||
clean_text = clean_text.replace(url, "").strip()
|
||||
|
||||
if clean_text:
|
||||
content.append({"type": "text", "text": clean_text})
|
||||
|
||||
for url in urls:
|
||||
try:
|
||||
data_uri = await _download_and_encode(url)
|
||||
content.append({"type": "image_url", "image_url": {"url": data_uri}})
|
||||
except Exception as e:
|
||||
logger.warning("Failed to download image %s: %s", url, e)
|
||||
# Put URL back in text
|
||||
if content and content[0].get("type") == "text":
|
||||
content[0]["text"] += f" {url}" # type: ignore[operator]
|
||||
else:
|
||||
content.insert(0, {"type": "text", "text": url})
|
||||
|
||||
return content
|
||||
|
||||
|
||||
async def process_messages(messages: list[ChatMessage]) -> list[ChatMessage]:
|
||||
"""Process all messages, downloading image URLs to base64.
|
||||
|
||||
This ensures uniform image format for all backends. Images arrive
|
||||
as base64 data URIs regardless of whether the client sent URLs or
|
||||
base64 originally.
|
||||
|
||||
Args:
|
||||
messages: Original messages (may contain image URLs).
|
||||
|
||||
Returns:
|
||||
Processed messages with images as base64 data URIs.
|
||||
"""
|
||||
processed: list[ChatMessage] = []
|
||||
|
||||
for msg in messages:
|
||||
content = msg.content
|
||||
|
||||
# Already multimodal array (OpenAI format)
|
||||
if isinstance(content, list):
|
||||
new_content = []
|
||||
for item in content:
|
||||
new_content.append(await _process_content_item(item))
|
||||
processed.append(
|
||||
ChatMessage(role=msg.role, content=new_content, name=msg.name)
|
||||
)
|
||||
|
||||
# Plain text — check for embedded image URLs
|
||||
elif isinstance(content, str):
|
||||
new_content = await _process_text_with_urls(content)
|
||||
processed.append(
|
||||
ChatMessage(role=msg.role, content=new_content, name=msg.name)
|
||||
)
|
||||
|
||||
else:
|
||||
processed.append(msg)
|
||||
|
||||
return processed
|
||||
135
ai_platform/modules/llm-inference/src/llm_inference/logging.py
Normal file
135
ai_platform/modules/llm-inference/src/llm_inference/logging.py
Normal file
|
|
@ -0,0 +1,135 @@
|
|||
"""Structured logging with request_id context propagation.
|
||||
|
||||
This module provides:
|
||||
- ContextVar-based request_id tracking
|
||||
- Structured JSON logging option
|
||||
- Request ID injection into all log records
|
||||
"""
|
||||
|
||||
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.
|
||||
|
||||
Args:
|
||||
record: The log record to modify.
|
||||
|
||||
Returns:
|
||||
bool: Always True (never filters out records).
|
||||
"""
|
||||
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.
|
||||
|
||||
Args:
|
||||
record: The log record to format.
|
||||
|
||||
Returns:
|
||||
str: JSON-formatted log string.
|
||||
"""
|
||||
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", "-"),
|
||||
}
|
||||
|
||||
# Add exception info if present
|
||||
if record.exc_info:
|
||||
log_data["exception"] = self.formatException(record.exc_info)
|
||||
|
||||
# Add extra fields
|
||||
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.
|
||||
"""
|
||||
# Get the root logger for our package
|
||||
root_logger = logging.getLogger("llm_inference")
|
||||
root_logger.setLevel(getattr(logging, level.upper(), logging.INFO))
|
||||
|
||||
# Remove existing handlers
|
||||
for handler in root_logger.handlers[:]:
|
||||
root_logger.removeHandler(handler)
|
||||
|
||||
# Create console handler
|
||||
handler = logging.StreamHandler(sys.stdout)
|
||||
handler.setLevel(logging.DEBUG)
|
||||
|
||||
# Add request ID filter
|
||||
handler.addFilter(RequestIdFilter())
|
||||
|
||||
# Set formatter based on preference
|
||||
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)
|
||||
|
||||
# Prevent propagation to root logger
|
||||
root_logger.propagate = False
|
||||
|
||||
|
||||
def get_logger(name: str) -> logging.Logger:
|
||||
"""Get a logger with the llm_inference prefix.
|
||||
|
||||
Args:
|
||||
name: Logger name (will be prefixed with 'llm_inference.').
|
||||
|
||||
Returns:
|
||||
logging.Logger: Configured logger instance.
|
||||
"""
|
||||
if name.startswith("llm_inference."):
|
||||
return logging.getLogger(name)
|
||||
return logging.getLogger(f"llm_inference.{name}")
|
||||
|
||||
|
||||
def set_request_id(request_id: str | None) -> None:
|
||||
"""Set the request ID for the current context.
|
||||
|
||||
Args:
|
||||
request_id: The request ID to set, or None to clear.
|
||||
"""
|
||||
request_id_ctx.set(request_id)
|
||||
|
||||
|
||||
def get_request_id() -> str | None:
|
||||
"""Get the request ID for the current context.
|
||||
|
||||
Returns:
|
||||
str | None: The current request ID, or None if not set.
|
||||
"""
|
||||
return request_id_ctx.get()
|
||||
|
|
@ -0,0 +1,63 @@
|
|||
"""Custom Prometheus metrics for LLM inference (offer §1.10).
|
||||
|
||||
Registered on the default prometheus_client registry, so they are exposed by the
|
||||
existing /metrics endpoint. Gives per-model latency percentiles (p50/p90/p99 via
|
||||
``histogram_quantile``), token throughput and request/error counts.
|
||||
|
||||
VRAM and loaded-model counts per backend are best scraped from the vLLM server's
|
||||
own /metrics (it exports GPU cache usage natively); this module covers the
|
||||
gateway-level signals.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from prometheus_client import Counter, Gauge, Histogram
|
||||
|
||||
LLM_LATENCY = Histogram(
|
||||
"llm_request_duration_seconds",
|
||||
"End-to-end completion latency at the gateway, per model and backend.",
|
||||
["model", "backend"],
|
||||
buckets=(0.05, 0.1, 0.25, 0.5, 1, 2, 5, 10, 20, 30, 60, 120),
|
||||
)
|
||||
|
||||
LLM_TOKENS = Counter(
|
||||
"llm_tokens_total",
|
||||
"Tokens processed, per model, backend and kind (prompt/completion).",
|
||||
["model", "backend", "kind"],
|
||||
)
|
||||
|
||||
LLM_REQUESTS = Counter(
|
||||
"llm_requests_total",
|
||||
"Completion requests, per model, backend and status (success/error).",
|
||||
["model", "backend", "status"],
|
||||
)
|
||||
|
||||
LLM_INFLIGHT = Gauge(
|
||||
"llm_inflight_requests",
|
||||
"In-flight completion requests, per backend.",
|
||||
["backend"],
|
||||
)
|
||||
|
||||
|
||||
def observe_success(
|
||||
model: str, backend: str, duration_s: float, usage: object | None
|
||||
) -> None:
|
||||
"""Record a successful completion's latency, tokens and request count."""
|
||||
|
||||
LLM_LATENCY.labels(model=model, backend=backend).observe(duration_s)
|
||||
LLM_REQUESTS.labels(model=model, backend=backend, status="success").inc()
|
||||
if usage is not None:
|
||||
prompt = getattr(usage, "prompt_tokens", 0) or 0
|
||||
completion = getattr(usage, "completion_tokens", 0) or 0
|
||||
if prompt:
|
||||
LLM_TOKENS.labels(model=model, backend=backend, kind="prompt").inc(prompt)
|
||||
if completion:
|
||||
LLM_TOKENS.labels(
|
||||
model=model, backend=backend, kind="completion"
|
||||
).inc(completion)
|
||||
|
||||
|
||||
def observe_error(model: str, backend: str) -> None:
|
||||
"""Record a failed completion request."""
|
||||
|
||||
LLM_REQUESTS.labels(model=model, backend=backend, status="error").inc()
|
||||
222
ai_platform/modules/llm-inference/src/llm_inference/retry.py
Normal file
222
ai_platform/modules/llm-inference/src/llm_inference/retry.py
Normal file
|
|
@ -0,0 +1,222 @@
|
|||
"""Retry utilities with exponential backoff for transient failures."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import random
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import TypeVar
|
||||
|
||||
import httpx
|
||||
|
||||
from llm_inference.exceptions import (
|
||||
CompletionError,
|
||||
LLMConnectionError,
|
||||
LLMRateLimitError,
|
||||
LLMTimeoutError,
|
||||
)
|
||||
|
||||
logger = logging.getLogger("llm_inference.retry")
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
# Import provider-specific exceptions conditionally
|
||||
try:
|
||||
import litellm.exceptions as litellm_exc
|
||||
|
||||
LITELLM_RETRYABLE = (
|
||||
litellm_exc.RateLimitError,
|
||||
litellm_exc.Timeout,
|
||||
litellm_exc.ServiceUnavailableError,
|
||||
)
|
||||
LITELLM_NON_RETRYABLE = (
|
||||
litellm_exc.AuthenticationError,
|
||||
litellm_exc.BadRequestError,
|
||||
)
|
||||
except ImportError:
|
||||
LITELLM_RETRYABLE = ()
|
||||
LITELLM_NON_RETRYABLE = ()
|
||||
|
||||
try:
|
||||
import openai
|
||||
|
||||
OPENAI_RETRYABLE = (
|
||||
openai.RateLimitError,
|
||||
openai.APIConnectionError,
|
||||
openai.APITimeoutError,
|
||||
)
|
||||
OPENAI_NON_RETRYABLE = (
|
||||
openai.AuthenticationError,
|
||||
openai.BadRequestError,
|
||||
openai.NotFoundError,
|
||||
)
|
||||
except ImportError:
|
||||
OPENAI_RETRYABLE = ()
|
||||
OPENAI_NON_RETRYABLE = ()
|
||||
|
||||
# Combine all retryable exceptions
|
||||
RETRYABLE_EXCEPTIONS: tuple[type[Exception], ...] = (
|
||||
*LITELLM_RETRYABLE,
|
||||
*OPENAI_RETRYABLE,
|
||||
httpx.TimeoutException,
|
||||
httpx.ConnectError,
|
||||
)
|
||||
|
||||
# Combine all non-retryable exceptions
|
||||
NON_RETRYABLE_EXCEPTIONS: tuple[type[Exception], ...] = (
|
||||
*LITELLM_NON_RETRYABLE,
|
||||
*OPENAI_NON_RETRYABLE,
|
||||
)
|
||||
|
||||
|
||||
def is_retryable_exception(exc: Exception) -> bool:
|
||||
"""Check if an exception is retryable.
|
||||
|
||||
Args:
|
||||
exc: Exception to check.
|
||||
|
||||
Returns:
|
||||
bool: True if the exception is a transient failure that can be retried.
|
||||
"""
|
||||
return isinstance(exc, RETRYABLE_EXCEPTIONS)
|
||||
|
||||
|
||||
def extract_retry_after(exc: Exception) -> float | None:
|
||||
"""Extract retry-after value from rate limit exceptions.
|
||||
|
||||
Args:
|
||||
exc: Exception that may contain retry-after header.
|
||||
|
||||
Returns:
|
||||
float | None: Retry-after value in seconds, or None if not available.
|
||||
"""
|
||||
if hasattr(exc, "response") and exc.response is not None:
|
||||
headers = getattr(exc.response, "headers", {})
|
||||
retry_after = headers.get("retry-after")
|
||||
if retry_after:
|
||||
try:
|
||||
return float(retry_after)
|
||||
except ValueError:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def translate_exception(
|
||||
exc: Exception,
|
||||
backend: str,
|
||||
model: str | None = None,
|
||||
) -> Exception:
|
||||
"""Translate provider-specific exceptions to our exception types.
|
||||
|
||||
Args:
|
||||
exc: Original exception from provider.
|
||||
backend: Backend name for error context.
|
||||
model: Model name for error context.
|
||||
|
||||
Returns:
|
||||
Exception: Translated exception of appropriate type.
|
||||
"""
|
||||
# Rate limit errors
|
||||
if LITELLM_RETRYABLE and isinstance(exc, litellm_exc.RateLimitError):
|
||||
retry_after = extract_retry_after(exc)
|
||||
return LLMRateLimitError(str(exc), backend=backend, retry_after=retry_after)
|
||||
|
||||
if OPENAI_RETRYABLE and isinstance(exc, openai.RateLimitError):
|
||||
retry_after = extract_retry_after(exc)
|
||||
return LLMRateLimitError(str(exc), backend=backend, retry_after=retry_after)
|
||||
|
||||
# Timeout errors
|
||||
if LITELLM_RETRYABLE and isinstance(exc, litellm_exc.Timeout):
|
||||
return LLMTimeoutError(str(exc), backend=backend)
|
||||
|
||||
if OPENAI_RETRYABLE and isinstance(exc, openai.APITimeoutError):
|
||||
return LLMTimeoutError(str(exc), backend=backend)
|
||||
|
||||
if isinstance(exc, httpx.TimeoutException):
|
||||
return LLMTimeoutError(str(exc), backend=backend)
|
||||
|
||||
# Connection errors
|
||||
if OPENAI_RETRYABLE and isinstance(exc, openai.APIConnectionError):
|
||||
return LLMConnectionError(backend, reason=str(exc))
|
||||
|
||||
if isinstance(exc, httpx.ConnectError):
|
||||
return LLMConnectionError(backend, reason=str(exc))
|
||||
|
||||
# Default to CompletionError
|
||||
return CompletionError(str(exc), backend=backend, model=model)
|
||||
|
||||
|
||||
async def retry_with_backoff(
|
||||
func: Callable[[], Awaitable[T]],
|
||||
max_retries: int,
|
||||
min_wait: float,
|
||||
max_wait: float,
|
||||
backend: str,
|
||||
model: str | None = None,
|
||||
) -> T:
|
||||
"""Execute an async function with exponential backoff retry.
|
||||
|
||||
Args:
|
||||
func: Async function to execute (takes no arguments).
|
||||
max_retries: Maximum number of retry attempts.
|
||||
min_wait: Minimum wait time between retries in seconds.
|
||||
max_wait: Maximum wait time between retries in seconds.
|
||||
backend: Backend name for error context.
|
||||
model: Model name for error context.
|
||||
|
||||
Returns:
|
||||
T: The result of the function.
|
||||
|
||||
Raises:
|
||||
LLMRateLimitError: On rate limit with no retries left.
|
||||
LLMTimeoutError: On timeout with no retries left.
|
||||
LLMConnectionError: On connection failure with no retries left.
|
||||
CompletionError: On other API errors.
|
||||
"""
|
||||
last_exception: Exception | None = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
return await func()
|
||||
|
||||
except NON_RETRYABLE_EXCEPTIONS as e:
|
||||
# Don't retry authentication or bad request errors
|
||||
raise translate_exception(e, backend, model) from e
|
||||
|
||||
except RETRYABLE_EXCEPTIONS as e:
|
||||
last_exception = e
|
||||
|
||||
if attempt >= max_retries:
|
||||
# No more retries
|
||||
raise translate_exception(e, backend, model) from e
|
||||
|
||||
# Calculate wait time with exponential backoff + jitter
|
||||
retry_after = extract_retry_after(e)
|
||||
if retry_after:
|
||||
# Respect the API's retry-after header, but cap at max_wait
|
||||
# to prevent malicious/buggy APIs from causing indefinite waits
|
||||
wait_time = min(retry_after, max_wait)
|
||||
else:
|
||||
# Exponential backoff: min_wait * 2^attempt (capped at max_wait)
|
||||
wait_time = min(min_wait * (2**attempt), max_wait)
|
||||
# Add jitter (0-25% of wait time)
|
||||
wait_time += random.uniform(0, wait_time * 0.25)
|
||||
|
||||
logger.warning(
|
||||
"Retry %d/%d for %s after %.1fs: %s",
|
||||
attempt + 1,
|
||||
max_retries,
|
||||
backend,
|
||||
wait_time,
|
||||
str(e),
|
||||
)
|
||||
|
||||
await asyncio.sleep(wait_time)
|
||||
|
||||
except Exception as e:
|
||||
# Unknown exception - translate and raise without retry
|
||||
raise translate_exception(e, backend, model) from e
|
||||
|
||||
# Should not reach here, but just in case
|
||||
if last_exception:
|
||||
raise translate_exception(last_exception, backend, model)
|
||||
raise CompletionError("Unknown error occurred", backend=backend, model=model)
|
||||
|
|
@ -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)
|
||||
223
ai_platform/modules/llm-inference/src/llm_inference/schemas.py
Normal file
223
ai_platform/modules/llm-inference/src/llm_inference/schemas.py
Normal file
|
|
@ -0,0 +1,223 @@
|
|||
"""Pydantic V2 API request/response schemas."""
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
|
||||
from llm_inference.types import (
|
||||
BackendType,
|
||||
ChatMessage,
|
||||
Choice,
|
||||
ModelInfo,
|
||||
StreamChoice,
|
||||
Usage,
|
||||
)
|
||||
|
||||
# =============================================================================
|
||||
# Completion Schemas
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class CompletionRequest(BaseModel):
|
||||
"""Request schema for chat completions."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
messages: list[ChatMessage] = Field(
|
||||
min_length=1,
|
||||
max_length=1000,
|
||||
description="List of messages in the conversation (1-1000 messages)",
|
||||
)
|
||||
model: str = Field(
|
||||
description="Model to use for completion",
|
||||
)
|
||||
temperature: float = Field(
|
||||
default=0.7,
|
||||
ge=0.0,
|
||||
le=2.0,
|
||||
description="Sampling temperature",
|
||||
)
|
||||
max_tokens: int | None = Field(
|
||||
default=None,
|
||||
ge=1,
|
||||
le=1000000,
|
||||
description="Maximum tokens to generate (1-1000000)",
|
||||
)
|
||||
stream: bool = Field(
|
||||
default=False,
|
||||
description="Enable streaming response",
|
||||
)
|
||||
backend: BackendType | None = Field(
|
||||
default=None,
|
||||
description="Backend to use (overrides default)",
|
||||
)
|
||||
top_p: float | None = Field(
|
||||
default=None,
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
description="Top-p sampling parameter",
|
||||
)
|
||||
frequency_penalty: float | None = Field(
|
||||
default=None,
|
||||
ge=-2.0,
|
||||
le=2.0,
|
||||
description="Frequency penalty",
|
||||
)
|
||||
presence_penalty: float | None = Field(
|
||||
default=None,
|
||||
ge=-2.0,
|
||||
le=2.0,
|
||||
description="Presence penalty",
|
||||
)
|
||||
stop: list[str] | str | None = Field(
|
||||
default=None,
|
||||
description="Stop sequences",
|
||||
)
|
||||
|
||||
@field_validator("messages")
|
||||
@classmethod
|
||||
def validate_messages_content(cls, v: list[ChatMessage]) -> list[ChatMessage]:
|
||||
"""Validate that messages have content (except assistant messages)."""
|
||||
for i, msg in enumerate(v):
|
||||
# Assistant messages can have empty content (for function calls etc)
|
||||
if msg.role != "assistant" and msg.content is None:
|
||||
raise ValueError(
|
||||
f"Message at index {i} with role '{msg.role}' cannot have empty content"
|
||||
)
|
||||
return v
|
||||
|
||||
|
||||
class CompletionResponse(BaseModel):
|
||||
"""Response schema for chat completions."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
id: str = Field(description="Unique completion ID")
|
||||
object: str = Field(default="chat.completion")
|
||||
created: int = Field(description="Unix timestamp of creation")
|
||||
model: str = Field(description="Model used for completion")
|
||||
choices: list[Choice] = Field(description="Completion choices")
|
||||
usage: Usage | None = Field(default=None, description="Token usage")
|
||||
backend: str = Field(description="Backend that served the request")
|
||||
|
||||
|
||||
class CompletionChunk(BaseModel):
|
||||
"""Streaming chunk for chat completions."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
id: str = Field(description="Unique completion ID")
|
||||
object: str = Field(default="chat.completion.chunk")
|
||||
created: int = Field(description="Unix timestamp of creation")
|
||||
model: str = Field(description="Model used for completion")
|
||||
choices: list[StreamChoice] = Field(description="Streaming choices")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Text Completion Schemas (legacy OpenAI /v1/completions)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TextChoice(BaseModel):
|
||||
"""One text-completion choice."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
index: int = Field(description="Choice index")
|
||||
text: str = Field(description="Generated text")
|
||||
finish_reason: str | None = Field(default=None, description="Finish reason")
|
||||
|
||||
|
||||
class TextCompletionRequest(BaseModel):
|
||||
"""Request schema for text completions (/v1/completions)."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
prompt: str | list[str] = Field(description="Prompt(s) to complete")
|
||||
model: str | None = Field(default=None, description="Model (alias) to use")
|
||||
temperature: float = Field(default=0.7, ge=0.0, le=2.0)
|
||||
max_tokens: int | None = Field(default=None, ge=1, le=1000000)
|
||||
backend: BackendType | None = Field(
|
||||
default=None, description="Backend to use (overrides default)"
|
||||
)
|
||||
top_p: float | None = Field(default=None, ge=0.0, le=1.0)
|
||||
frequency_penalty: float | None = Field(default=None, ge=-2.0, le=2.0)
|
||||
presence_penalty: float | None = Field(default=None, ge=-2.0, le=2.0)
|
||||
stop: list[str] | str | None = Field(default=None, description="Stop sequences")
|
||||
|
||||
|
||||
class TextCompletionResponse(BaseModel):
|
||||
"""Response schema for text completions."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
id: str = Field(description="Unique completion ID")
|
||||
object: str = Field(default="text_completion")
|
||||
created: int = Field(description="Unix timestamp of creation")
|
||||
model: str = Field(description="Model used for completion")
|
||||
choices: list[TextChoice] = Field(description="Completion choices")
|
||||
usage: Usage | None = Field(default=None, description="Token usage")
|
||||
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")
|
||||
|
||||
|
||||
class ModelLoadRequest(BaseModel):
|
||||
"""Request schema for loading a model."""
|
||||
|
||||
model: str = Field(description="Model identifier to load")
|
||||
backend: BackendType = Field(description="Backend to load the model on")
|
||||
|
||||
|
||||
class ModelLoadResponse(BaseModel):
|
||||
"""Response schema for model load/unload operations."""
|
||||
|
||||
success: bool = Field(description="Whether the operation succeeded")
|
||||
model: str = Field(description="Model identifier")
|
||||
backend: str = Field(description="Backend the operation was performed on")
|
||||
message: str | None = Field(default=None, description="Optional message")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# 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")
|
||||
87
ai_platform/modules/llm-inference/src/llm_inference/types.py
Normal file
87
ai_platform/modules/llm-inference/src/llm_inference/types.py
Normal file
|
|
@ -0,0 +1,87 @@
|
|||
"""Core types and enums for LLM inference."""
|
||||
|
||||
from enum import Enum
|
||||
from typing import Literal
|
||||
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
class BackendType(str, Enum):
|
||||
"""Supported LLM backends."""
|
||||
|
||||
LITELLM = "litellm"
|
||||
VLLM = "vllm"
|
||||
LLAMACPP = "llamacpp"
|
||||
|
||||
|
||||
class ChatMessage(BaseModel):
|
||||
"""A chat message in a conversation."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
role: Literal["system", "user", "assistant", "function", "tool"] = Field(
|
||||
description="The role of the message author",
|
||||
)
|
||||
content: str | list[dict[str, Any]] | None = Field(
|
||||
default=None,
|
||||
description="The content of the message (string or multimodal content array)",
|
||||
)
|
||||
name: str | None = Field(
|
||||
default=None,
|
||||
description="Optional name for the message author",
|
||||
)
|
||||
|
||||
|
||||
class Usage(BaseModel):
|
||||
"""Token usage information."""
|
||||
|
||||
prompt_tokens: int = Field(default=0)
|
||||
completion_tokens: int = Field(default=0)
|
||||
total_tokens: int = Field(default=0)
|
||||
|
||||
|
||||
class Choice(BaseModel):
|
||||
"""A completion choice."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
index: int = Field(default=0)
|
||||
message: ChatMessage | None = Field(default=None)
|
||||
finish_reason: str | None = Field(default=None)
|
||||
|
||||
|
||||
class Delta(BaseModel):
|
||||
"""Streaming delta content."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
role: str | None = Field(default=None)
|
||||
content: str | None = Field(default=None)
|
||||
|
||||
|
||||
class StreamChoice(BaseModel):
|
||||
"""A streaming completion choice."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
index: int = Field(default=0)
|
||||
delta: Delta = Field(default_factory=Delta)
|
||||
finish_reason: str | None = Field(default=None)
|
||||
|
||||
|
||||
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")
|
||||
context_length: int | None = Field(
|
||||
default=None,
|
||||
description="Maximum context length in tokens",
|
||||
)
|
||||
capabilities: list[str] = Field(
|
||||
default_factory=list,
|
||||
description="Model capabilities (e.g., 'chat', 'completion')",
|
||||
)
|
||||
29
ai_platform/modules/llm-inference/src/llm_inference/utils.py
Normal file
29
ai_platform/modules/llm-inference/src/llm_inference/utils.py
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
"""Shared utilities for LLM inference module."""
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
|
||||
async def safe_close_stream(
|
||||
stream: Any | None,
|
||||
logger: logging.Logger,
|
||||
) -> None:
|
||||
"""Safely close an async stream, handling both aclose() and close().
|
||||
|
||||
Attempts aclose() first (async close), falls back to close() if not available.
|
||||
Silently handles any exceptions during cleanup.
|
||||
|
||||
Args:
|
||||
stream: Stream object to close, or None.
|
||||
logger: Logger for debug messages on cleanup errors.
|
||||
"""
|
||||
if stream is None:
|
||||
return
|
||||
|
||||
try:
|
||||
if hasattr(stream, "aclose"):
|
||||
await stream.aclose()
|
||||
elif hasattr(stream, "close"):
|
||||
await stream.close()
|
||||
except Exception as e:
|
||||
logger.debug("Error during stream cleanup: %s", e)
|
||||
1
ai_platform/modules/llm-inference/tests/__init__.py
Normal file
1
ai_platform/modules/llm-inference/tests/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
"""Tests for LLM inference module."""
|
||||
162
ai_platform/modules/llm-inference/tests/conftest.py
Normal file
162
ai_platform/modules/llm-inference/tests/conftest.py
Normal file
|
|
@ -0,0 +1,162 @@
|
|||
"""Pytest fixtures for LLM inference tests."""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from llm_inference.api.app import create_app
|
||||
from llm_inference.api.dependencies import init_concurrency_limiter
|
||||
from llm_inference.client import LLMClient
|
||||
from llm_inference.config import LLMSettings, SettingsCache
|
||||
from llm_inference.schemas import CompletionChunk, CompletionResponse
|
||||
from llm_inference.types import (
|
||||
BackendType,
|
||||
ChatMessage,
|
||||
Choice,
|
||||
Delta,
|
||||
StreamChoice,
|
||||
Usage,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clear_settings_cache() -> None:
|
||||
"""Clear settings cache before each test."""
|
||||
SettingsCache.clear()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_settings() -> LLMSettings:
|
||||
"""Create test settings with mocked values."""
|
||||
return LLMSettings(
|
||||
default_backend="litellm",
|
||||
default_model="test-model",
|
||||
openrouter_api_key="test-openrouter-key",
|
||||
openai_api_key="test-openai-key",
|
||||
enable_vllm=False,
|
||||
enable_llamacpp=False,
|
||||
host="127.0.0.1",
|
||||
port=8100,
|
||||
external_url="http://localhost:8100",
|
||||
api_tokens=None, # Auth disabled by default in tests
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_completion_response() -> CompletionResponse:
|
||||
"""Create a mock completion response."""
|
||||
return CompletionResponse(
|
||||
id="test-completion-id",
|
||||
created=1234567890,
|
||||
model="test-model",
|
||||
choices=[
|
||||
Choice(
|
||||
index=0,
|
||||
message=ChatMessage(role="assistant", content="Hello! How can I help?"),
|
||||
finish_reason="stop",
|
||||
)
|
||||
],
|
||||
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
||||
backend="litellm",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_completion_chunks() -> list[CompletionChunk]:
|
||||
"""Create mock streaming completion chunks."""
|
||||
return [
|
||||
CompletionChunk(
|
||||
id="test-chunk-1",
|
||||
created=1234567890,
|
||||
model="test-model",
|
||||
choices=[
|
||||
StreamChoice(
|
||||
index=0,
|
||||
delta=Delta(role="assistant", content="Hello"),
|
||||
finish_reason=None,
|
||||
)
|
||||
],
|
||||
),
|
||||
CompletionChunk(
|
||||
id="test-chunk-2",
|
||||
created=1234567891,
|
||||
model="test-model",
|
||||
choices=[
|
||||
StreamChoice(
|
||||
index=0,
|
||||
delta=Delta(content="!"),
|
||||
finish_reason=None,
|
||||
)
|
||||
],
|
||||
),
|
||||
CompletionChunk(
|
||||
id="test-chunk-3",
|
||||
created=1234567892,
|
||||
model="test-model",
|
||||
choices=[
|
||||
StreamChoice(
|
||||
index=0,
|
||||
delta=Delta(),
|
||||
finish_reason="stop",
|
||||
)
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_litellm_response() -> MagicMock:
|
||||
"""Create a mock LiteLLM response object."""
|
||||
response = MagicMock()
|
||||
response.id = "test-completion-id"
|
||||
response.created = 1234567890
|
||||
response.model = "test-model"
|
||||
|
||||
choice = MagicMock()
|
||||
choice.index = 0
|
||||
choice.message.role = "assistant"
|
||||
choice.message.content = "Hello! How can I help?"
|
||||
choice.finish_reason = "stop"
|
||||
response.choices = [choice]
|
||||
|
||||
response.usage.prompt_tokens = 10
|
||||
response.usage.completion_tokens = 5
|
||||
response.usage.total_tokens = 15
|
||||
|
||||
return response
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client_with_mock_backend(
|
||||
test_settings: LLMSettings, mock_completion_response: CompletionResponse
|
||||
) -> LLMClient:
|
||||
"""Create an LLMClient with mocked backend."""
|
||||
# Inject test settings into cache
|
||||
SettingsCache.set(test_settings)
|
||||
|
||||
client = LLMClient(settings=test_settings)
|
||||
|
||||
# Mock the backend's complete method
|
||||
backend = client.registry.get(BackendType.LITELLM)
|
||||
backend.complete = AsyncMock(return_value=mock_completion_response)
|
||||
|
||||
return client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app_client(test_settings: LLMSettings) -> TestClient:
|
||||
"""Create a FastAPI TestClient with test settings."""
|
||||
# Inject test settings into cache BEFORE creating the app
|
||||
SettingsCache.set(test_settings)
|
||||
|
||||
# Initialize concurrency limiter for tests
|
||||
init_concurrency_limiter(test_settings.max_concurrent_completions)
|
||||
|
||||
app = create_app()
|
||||
|
||||
# Also set in app.state for handlers that access it
|
||||
app.state.settings = test_settings
|
||||
app.state.client = LLMClient(settings=test_settings)
|
||||
|
||||
return TestClient(app)
|
||||
209
ai_platform/modules/llm-inference/tests/test_api.py
Normal file
209
ai_platform/modules/llm-inference/tests/test_api.py
Normal file
|
|
@ -0,0 +1,209 @@
|
|||
"""Tests for FastAPI routes."""
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from llm_inference.api.app import create_app
|
||||
from llm_inference.api.dependencies import init_concurrency_limiter
|
||||
from llm_inference.client import LLMClient
|
||||
from llm_inference.config import LLMSettings, SettingsCache
|
||||
from llm_inference.schemas import CompletionResponse
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app_client(test_settings: LLMSettings) -> TestClient:
|
||||
"""Create a test client with mocked client."""
|
||||
# Inject test settings into cache BEFORE creating app
|
||||
SettingsCache.set(test_settings)
|
||||
|
||||
# Initialize concurrency limiter for tests
|
||||
init_concurrency_limiter(test_settings.max_concurrent_completions)
|
||||
|
||||
app = create_app()
|
||||
app.state.settings = test_settings
|
||||
app.state.client = LLMClient(settings=test_settings)
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
class TestHealthRoutes:
|
||||
"""Tests for health check routes."""
|
||||
|
||||
def test_health_endpoint(self, app_client: TestClient) -> None:
|
||||
"""Test /health endpoint."""
|
||||
response = app_client.get("/health")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "status" in data
|
||||
assert "backends" in data
|
||||
assert data["status"] in ["healthy", "degraded", "unhealthy"]
|
||||
|
||||
def test_ready_endpoint(self, app_client: TestClient) -> None:
|
||||
"""Test /ready endpoint."""
|
||||
response = app_client.get("/ready")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["ready"] is True
|
||||
|
||||
|
||||
class TestModelsRoutes:
|
||||
"""Tests for model management routes."""
|
||||
|
||||
def test_list_models(self, app_client: TestClient) -> None:
|
||||
"""Test GET /v1/models."""
|
||||
response = app_client.get("/v1/models")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "data" in data
|
||||
assert "object" in data
|
||||
assert data["object"] == "list"
|
||||
|
||||
def test_list_models_with_backend_filter(self, app_client: TestClient) -> None:
|
||||
"""Test GET /v1/models with backend filter."""
|
||||
response = app_client.get("/v1/models?backend=litellm")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
# All models should be from litellm backend
|
||||
for model in data["data"]:
|
||||
assert model["backend"] == "litellm"
|
||||
|
||||
def test_list_models_invalid_backend(self, app_client: TestClient) -> None:
|
||||
"""Test GET /v1/models with invalid backend."""
|
||||
response = app_client.get("/v1/models?backend=invalid")
|
||||
|
||||
assert response.status_code == 400
|
||||
|
||||
def test_list_backends(self, app_client: TestClient) -> None:
|
||||
"""Test GET /v1/backends."""
|
||||
response = app_client.get("/v1/backends")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "backends" in data
|
||||
assert "litellm" in data["backends"]
|
||||
|
||||
|
||||
class TestCompletionsRoutes:
|
||||
"""Tests for completion routes."""
|
||||
|
||||
def test_chat_completions(
|
||||
self,
|
||||
app_client: TestClient,
|
||||
mock_completion_response: CompletionResponse,
|
||||
) -> None:
|
||||
"""Test POST /v1/chat/completions."""
|
||||
# Mock the client's complete method
|
||||
with patch.object(
|
||||
app_client.app.state.client,
|
||||
"complete",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_completion_response,
|
||||
):
|
||||
response = app_client.post(
|
||||
"/v1/chat/completions",
|
||||
json={
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"model": "gpt-3.5-turbo",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "id" in data
|
||||
assert "choices" in data
|
||||
assert "model" in data
|
||||
|
||||
def test_chat_completions_with_parameters(
|
||||
self,
|
||||
app_client: TestClient,
|
||||
mock_completion_response: CompletionResponse,
|
||||
) -> None:
|
||||
"""Test completion with optional parameters."""
|
||||
with patch.object(
|
||||
app_client.app.state.client,
|
||||
"complete",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_completion_response,
|
||||
):
|
||||
response = app_client.post(
|
||||
"/v1/chat/completions",
|
||||
json={
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"model": "gpt-4",
|
||||
"temperature": 0.5,
|
||||
"max_tokens": 100,
|
||||
"top_p": 0.9,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
def test_chat_completions_with_backend_override(
|
||||
self,
|
||||
app_client: TestClient,
|
||||
mock_completion_response: CompletionResponse,
|
||||
) -> None:
|
||||
"""Test completion with backend override."""
|
||||
with patch.object(
|
||||
app_client.app.state.client,
|
||||
"complete",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_completion_response,
|
||||
):
|
||||
response = app_client.post(
|
||||
"/v1/chat/completions",
|
||||
json={
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"model": "gpt-3.5-turbo",
|
||||
"backend": "litellm",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
def test_chat_completions_missing_messages(self, app_client: TestClient) -> None:
|
||||
"""Test completion without messages returns error."""
|
||||
response = app_client.post(
|
||||
"/v1/chat/completions",
|
||||
json={"model": "gpt-3.5-turbo"},
|
||||
)
|
||||
|
||||
assert response.status_code == 422 # Validation error
|
||||
|
||||
def test_chat_completions_missing_model(self, app_client: TestClient) -> None:
|
||||
"""Test completion without model returns error."""
|
||||
response = app_client.post(
|
||||
"/v1/chat/completions",
|
||||
json={"messages": [{"role": "user", "content": "Hello"}]},
|
||||
)
|
||||
|
||||
assert response.status_code == 422 # Validation error
|
||||
|
||||
|
||||
class TestModelLoadRoutes:
|
||||
"""Tests for model load/unload routes."""
|
||||
|
||||
def test_load_model_unsupported_backend(self, app_client: TestClient) -> None:
|
||||
"""Test loading model on backend that doesn't support it."""
|
||||
response = app_client.post(
|
||||
"/v1/models/load",
|
||||
json={"model": "test-model", "backend": "litellm"},
|
||||
)
|
||||
|
||||
# LiteLLM doesn't support model loading
|
||||
assert response.status_code == 400
|
||||
|
||||
def test_unload_model_unsupported_backend(self, app_client: TestClient) -> None:
|
||||
"""Test unloading model on backend that doesn't support it."""
|
||||
response = app_client.post(
|
||||
"/v1/models/unload",
|
||||
json={"model": "test-model", "backend": "litellm"},
|
||||
)
|
||||
|
||||
# LiteLLM doesn't support model unloading
|
||||
assert response.status_code == 400
|
||||
224
ai_platform/modules/llm-inference/tests/test_auth.py
Normal file
224
ai_platform/modules/llm-inference/tests/test_auth.py
Normal file
|
|
@ -0,0 +1,224 @@
|
|||
"""Tests for API authentication."""
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from llm_inference.api.app import create_app
|
||||
from llm_inference.api.dependencies import init_concurrency_limiter
|
||||
from llm_inference.client import LLMClient
|
||||
from llm_inference.config import LLMSettings, SettingsCache
|
||||
|
||||
|
||||
class TestAuthDisabled:
|
||||
"""Tests when authentication is disabled (no tokens configured)."""
|
||||
|
||||
@pytest.fixture
|
||||
def app_client_no_auth(self, test_settings: LLMSettings) -> TestClient:
|
||||
"""Create test client with auth disabled."""
|
||||
# test_settings has api_tokens=None by default
|
||||
SettingsCache.set(test_settings)
|
||||
init_concurrency_limiter(test_settings.max_concurrent_completions)
|
||||
app = create_app()
|
||||
app.state.settings = test_settings
|
||||
app.state.client = LLMClient(settings=test_settings)
|
||||
return TestClient(app)
|
||||
|
||||
def test_models_accessible_without_token(
|
||||
self, app_client_no_auth: TestClient
|
||||
) -> None:
|
||||
"""Test that models endpoint works without token when auth disabled."""
|
||||
response = app_client_no_auth.get("/v1/models")
|
||||
assert response.status_code != 401
|
||||
|
||||
def test_backends_accessible_without_token(
|
||||
self, app_client_no_auth: TestClient
|
||||
) -> None:
|
||||
"""Test that backends endpoint works without token when auth disabled."""
|
||||
response = app_client_no_auth.get("/v1/backends")
|
||||
assert response.status_code == 200
|
||||
|
||||
def test_health_accessible(self, app_client_no_auth: TestClient) -> None:
|
||||
"""Test health endpoint accessible."""
|
||||
response = app_client_no_auth.get("/health")
|
||||
assert response.status_code == 200
|
||||
|
||||
def test_ready_accessible(self, app_client_no_auth: TestClient) -> None:
|
||||
"""Test ready endpoint accessible."""
|
||||
response = app_client_no_auth.get("/ready")
|
||||
# May return 200 or 503 depending on backend health, but not 401
|
||||
assert response.status_code != 401
|
||||
|
||||
|
||||
class TestAuthEnabled:
|
||||
"""Tests when authentication is enabled."""
|
||||
|
||||
@pytest.fixture
|
||||
def auth_settings(self, test_settings: LLMSettings) -> LLMSettings:
|
||||
"""Create settings with auth enabled."""
|
||||
return LLMSettings(
|
||||
default_backend=test_settings.default_backend,
|
||||
enable_vllm=test_settings.enable_vllm,
|
||||
enable_llamacpp=test_settings.enable_llamacpp,
|
||||
external_url=test_settings.external_url,
|
||||
openrouter_api_key=test_settings.openrouter_api_key,
|
||||
api_tokens=frozenset(["test-token-1", "test-token-2"]),
|
||||
)
|
||||
|
||||
@pytest.fixture
|
||||
def app_client_auth(self, auth_settings: LLMSettings) -> TestClient:
|
||||
"""Create test client with auth enabled."""
|
||||
SettingsCache.set(auth_settings)
|
||||
init_concurrency_limiter(auth_settings.max_concurrent_completions)
|
||||
app = create_app()
|
||||
app.state.settings = auth_settings
|
||||
app.state.client = LLMClient(settings=auth_settings)
|
||||
return TestClient(app)
|
||||
|
||||
def test_missing_token_returns_401(self, app_client_auth: TestClient) -> None:
|
||||
"""Test 401 returned when token missing."""
|
||||
response = app_client_auth.get("/v1/models")
|
||||
assert response.status_code == 401
|
||||
assert "WWW-Authenticate" in response.headers
|
||||
assert response.headers["WWW-Authenticate"] == "Bearer"
|
||||
data = response.json()
|
||||
assert data["detail"]["error"] == "Authentication required"
|
||||
|
||||
def test_invalid_token_returns_401(self, app_client_auth: TestClient) -> None:
|
||||
"""Test 401 returned for invalid token."""
|
||||
response = app_client_auth.get(
|
||||
"/v1/models",
|
||||
headers={"Authorization": "Bearer invalid-token"},
|
||||
)
|
||||
assert response.status_code == 401
|
||||
data = response.json()
|
||||
assert data["detail"]["error"] == "Authentication failed"
|
||||
|
||||
def test_valid_token_allows_access(self, app_client_auth: TestClient) -> None:
|
||||
"""Test valid token grants access."""
|
||||
response = app_client_auth.get(
|
||||
"/v1/models",
|
||||
headers={"Authorization": "Bearer test-token-1"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
def test_second_valid_token_works(self, app_client_auth: TestClient) -> None:
|
||||
"""Test second token also works."""
|
||||
response = app_client_auth.get(
|
||||
"/v1/models",
|
||||
headers={"Authorization": "Bearer test-token-2"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
def test_malformed_header_missing_bearer(self, app_client_auth: TestClient) -> None:
|
||||
"""Test malformed Authorization header without Bearer prefix."""
|
||||
response = app_client_auth.get(
|
||||
"/v1/models",
|
||||
headers={"Authorization": "test-token-1"},
|
||||
)
|
||||
assert response.status_code == 401
|
||||
|
||||
def test_malformed_header_wrong_prefix(self, app_client_auth: TestClient) -> None:
|
||||
"""Test Authorization header with wrong prefix."""
|
||||
response = app_client_auth.get(
|
||||
"/v1/models",
|
||||
headers={"Authorization": "Basic test-token-1"},
|
||||
)
|
||||
assert response.status_code == 401
|
||||
|
||||
def test_health_excluded_from_auth(self, app_client_auth: TestClient) -> None:
|
||||
"""Test /health accessible without token even when auth enabled."""
|
||||
response = app_client_auth.get("/health")
|
||||
assert response.status_code == 200
|
||||
|
||||
def test_ready_excluded_from_auth(self, app_client_auth: TestClient) -> None:
|
||||
"""Test /ready accessible without token even when auth enabled."""
|
||||
response = app_client_auth.get("/ready")
|
||||
# May return 200 or 503 depending on backend health, but not 401
|
||||
assert response.status_code != 401
|
||||
|
||||
def test_backends_requires_auth(self, app_client_auth: TestClient) -> None:
|
||||
"""Test /v1/backends requires authentication."""
|
||||
response = app_client_auth.get("/v1/backends")
|
||||
assert response.status_code == 401
|
||||
|
||||
def test_backends_with_valid_token(self, app_client_auth: TestClient) -> None:
|
||||
"""Test /v1/backends works with valid token."""
|
||||
response = app_client_auth.get(
|
||||
"/v1/backends",
|
||||
headers={"Authorization": "Bearer test-token-1"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
class TestTokenParsing:
|
||||
"""Tests for token configuration parsing."""
|
||||
|
||||
def test_comma_separated_tokens(self) -> None:
|
||||
"""Test parsing comma-separated tokens."""
|
||||
settings = LLMSettings(
|
||||
default_backend="litellm",
|
||||
enable_vllm=False,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:8100",
|
||||
api_tokens="token1,token2,token3",
|
||||
)
|
||||
assert settings.api_tokens == frozenset(["token1", "token2", "token3"])
|
||||
assert settings.auth_enabled is True
|
||||
|
||||
def test_empty_string_disables_auth(self) -> None:
|
||||
"""Test empty string results in auth disabled."""
|
||||
settings = LLMSettings(
|
||||
default_backend="litellm",
|
||||
enable_vllm=False,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:8100",
|
||||
api_tokens="",
|
||||
)
|
||||
assert settings.api_tokens is None
|
||||
assert settings.auth_enabled is False
|
||||
|
||||
def test_whitespace_tokens_stripped(self) -> None:
|
||||
"""Test whitespace around tokens is stripped."""
|
||||
settings = LLMSettings(
|
||||
default_backend="litellm",
|
||||
enable_vllm=False,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:8100",
|
||||
api_tokens=" token1 , token2 , token3 ",
|
||||
)
|
||||
assert settings.api_tokens == frozenset(["token1", "token2", "token3"])
|
||||
|
||||
def test_none_disables_auth(self) -> None:
|
||||
"""Test None results in auth disabled."""
|
||||
settings = LLMSettings(
|
||||
default_backend="litellm",
|
||||
enable_vllm=False,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:8100",
|
||||
)
|
||||
assert settings.api_tokens is None
|
||||
assert settings.auth_enabled is False
|
||||
|
||||
def test_single_token(self) -> None:
|
||||
"""Test single token without comma."""
|
||||
settings = LLMSettings(
|
||||
default_backend="litellm",
|
||||
enable_vllm=False,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:8100",
|
||||
api_tokens="single-token",
|
||||
)
|
||||
assert settings.api_tokens == frozenset(["single-token"])
|
||||
assert settings.auth_enabled is True
|
||||
|
||||
def test_whitespace_only_disables_auth(self) -> None:
|
||||
"""Test whitespace-only string results in auth disabled."""
|
||||
settings = LLMSettings(
|
||||
default_backend="litellm",
|
||||
enable_vllm=False,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:8100",
|
||||
api_tokens=" , , ",
|
||||
)
|
||||
assert settings.api_tokens is None
|
||||
assert settings.auth_enabled is False
|
||||
201
ai_platform/modules/llm-inference/tests/test_client.py
Normal file
201
ai_platform/modules/llm-inference/tests/test_client.py
Normal file
|
|
@ -0,0 +1,201 @@
|
|||
"""Tests for LLMClient."""
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from llm_inference.client import LLMClient
|
||||
from llm_inference.config import LLMSettings, SettingsCache
|
||||
from llm_inference.exceptions import BackendNotEnabledError
|
||||
from llm_inference.schemas import CompletionResponse
|
||||
from llm_inference.types import BackendType, ChatMessage
|
||||
|
||||
|
||||
class TestLLMClient:
|
||||
"""Tests for LLMClient."""
|
||||
|
||||
def test_init_with_settings(self, test_settings: LLMSettings) -> None:
|
||||
"""Test client initialization with custom settings."""
|
||||
client = LLMClient(settings=test_settings)
|
||||
|
||||
assert client._settings == test_settings
|
||||
assert client.registry is not None
|
||||
|
||||
def test_init_without_settings(self, test_settings: LLMSettings) -> None:
|
||||
"""Test client initialization uses cached settings."""
|
||||
# Inject settings into cache
|
||||
SettingsCache.set(test_settings)
|
||||
|
||||
client = LLMClient()
|
||||
assert client._settings is not None
|
||||
assert client._settings.default_backend == test_settings.default_backend
|
||||
|
||||
def test_list_backends(self, test_settings: LLMSettings) -> None:
|
||||
"""Test listing available backends."""
|
||||
client = LLMClient(settings=test_settings)
|
||||
backends = client.list_backends()
|
||||
|
||||
# LiteLLM should always be available
|
||||
assert BackendType.LITELLM in backends
|
||||
|
||||
# vLLM and llamacpp should not be available (disabled)
|
||||
assert BackendType.VLLM not in backends
|
||||
assert BackendType.LLAMACPP not in backends
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_complete_uses_default_backend(
|
||||
self, test_settings: LLMSettings, mock_completion_response: CompletionResponse
|
||||
) -> None:
|
||||
"""Test that complete uses default backend when not specified."""
|
||||
client = LLMClient(settings=test_settings)
|
||||
|
||||
# Mock the backend's complete method
|
||||
with patch.object(
|
||||
client.registry.get(BackendType.LITELLM),
|
||||
"complete",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_completion_response,
|
||||
) as mock_complete:
|
||||
response = await client.complete(
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
)
|
||||
|
||||
assert response.backend == "litellm"
|
||||
mock_complete.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_complete_uses_default_model(
|
||||
self, test_settings: LLMSettings, mock_completion_response: CompletionResponse
|
||||
) -> None:
|
||||
"""Test that complete uses default model when not specified."""
|
||||
client = LLMClient(settings=test_settings)
|
||||
|
||||
with patch.object(
|
||||
client.registry.get(BackendType.LITELLM),
|
||||
"complete",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_completion_response,
|
||||
) as mock_complete:
|
||||
await client.complete(
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
)
|
||||
|
||||
# Check that default model was used
|
||||
call_args = mock_complete.call_args
|
||||
assert call_args[0][1] == test_settings.default_model
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_complete_with_dict_messages(
|
||||
self, test_settings: LLMSettings, mock_completion_response: CompletionResponse
|
||||
) -> None:
|
||||
"""Test that complete accepts dict messages."""
|
||||
client = LLMClient(settings=test_settings)
|
||||
|
||||
with patch.object(
|
||||
client.registry.get(BackendType.LITELLM),
|
||||
"complete",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_completion_response,
|
||||
):
|
||||
response = await client.complete(
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
model="gpt-4",
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_complete_with_chatmessage_objects(
|
||||
self, test_settings: LLMSettings, mock_completion_response: CompletionResponse
|
||||
) -> None:
|
||||
"""Test that complete accepts ChatMessage objects."""
|
||||
client = LLMClient(settings=test_settings)
|
||||
|
||||
with patch.object(
|
||||
client.registry.get(BackendType.LITELLM),
|
||||
"complete",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_completion_response,
|
||||
):
|
||||
response = await client.complete(
|
||||
messages=[ChatMessage(role="user", content="Hello")],
|
||||
model="gpt-4",
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_models(self, test_settings: LLMSettings) -> None:
|
||||
"""Test listing models from all backends."""
|
||||
# Set API keys so models are returned
|
||||
test_settings.openai_api_key = "test-key"
|
||||
client = LLMClient(settings=test_settings)
|
||||
|
||||
models = await client.list_models()
|
||||
|
||||
# Should return some models
|
||||
assert len(models) > 0
|
||||
# All should be from litellm backend
|
||||
assert all(m.backend == "litellm" for m in models)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_check(self, test_settings: LLMSettings) -> None:
|
||||
"""Test health check returns backend status."""
|
||||
client = LLMClient(settings=test_settings)
|
||||
|
||||
health = await client.health_check()
|
||||
|
||||
# Should return status for litellm backend
|
||||
assert "litellm" in health
|
||||
# Health check may return True or False depending on connectivity
|
||||
assert isinstance(health["litellm"], bool)
|
||||
|
||||
def test_parse_messages_dict(self, test_settings: LLMSettings) -> None:
|
||||
"""Test message parsing from dicts."""
|
||||
client = LLMClient(settings=test_settings)
|
||||
|
||||
messages = client._parse_messages([{"role": "user", "content": "Hello"}])
|
||||
|
||||
assert len(messages) == 1
|
||||
assert isinstance(messages[0], ChatMessage)
|
||||
assert messages[0].role == "user"
|
||||
assert messages[0].content == "Hello"
|
||||
|
||||
def test_parse_messages_chatmessage(self, test_settings: LLMSettings) -> None:
|
||||
"""Test message parsing from ChatMessage objects."""
|
||||
client = LLMClient(settings=test_settings)
|
||||
original = ChatMessage(role="user", content="Hello")
|
||||
|
||||
messages = client._parse_messages([original])
|
||||
|
||||
assert len(messages) == 1
|
||||
assert messages[0] is original
|
||||
|
||||
|
||||
class TestBackendSelection:
|
||||
"""Tests for backend selection."""
|
||||
|
||||
def test_unavailable_backend_raises_error(self, test_settings: LLMSettings) -> None:
|
||||
"""Test that requesting unavailable backend raises error."""
|
||||
client = LLMClient(settings=test_settings)
|
||||
|
||||
with pytest.raises(BackendNotEnabledError):
|
||||
client.registry.get(BackendType.VLLM)
|
||||
|
||||
def test_enabled_vllm_backend(self) -> None:
|
||||
"""Test that vLLM backend is available when enabled."""
|
||||
settings = LLMSettings(
|
||||
default_backend="litellm",
|
||||
enable_vllm=True,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:8100",
|
||||
)
|
||||
|
||||
# This will fail because openai package may not be installed
|
||||
# but we can at least verify the registry tries to load it
|
||||
client = LLMClient(settings=settings)
|
||||
|
||||
# vLLM might be in backends if openai package is available
|
||||
backends = client.list_backends()
|
||||
# Just verify litellm is there
|
||||
assert BackendType.LITELLM in backends
|
||||
139
ai_platform/modules/llm-inference/tests/test_concurrency.py
Normal file
139
ai_platform/modules/llm-inference/tests/test_concurrency.py
Normal file
|
|
@ -0,0 +1,139 @@
|
|||
"""Tests for concurrency limiter functionality."""
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from llm_inference.api.dependencies import ConcurrencyLimiter
|
||||
|
||||
|
||||
class TestConcurrencyLimiter:
|
||||
"""Tests for ConcurrencyLimiter class."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acquire_within_limit(self) -> None:
|
||||
"""Test acquiring slots within concurrency limit."""
|
||||
limiter = ConcurrencyLimiter(max_concurrent=2)
|
||||
|
||||
async with limiter.acquire():
|
||||
assert limiter.current_count == 1
|
||||
assert limiter.available == 1
|
||||
|
||||
assert limiter.current_count == 0
|
||||
assert limiter.available == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acquire_exceeds_limit_raises_503(self) -> None:
|
||||
"""Test that exceeding limit raises 503 HTTPException."""
|
||||
limiter = ConcurrencyLimiter(max_concurrent=1)
|
||||
|
||||
async with limiter.acquire():
|
||||
# Try to acquire another slot while one is held
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
async with limiter.acquire():
|
||||
pass
|
||||
|
||||
assert exc_info.value.status_code == 503
|
||||
assert "Too many concurrent requests" in str(exc_info.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acquire_blocking_mode(self) -> None:
|
||||
"""Test blocking mode waits for available slot."""
|
||||
limiter = ConcurrencyLimiter(max_concurrent=1)
|
||||
results: list[int] = []
|
||||
|
||||
async def task(task_id: int) -> None:
|
||||
async with limiter.acquire(blocking=True):
|
||||
results.append(task_id)
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
# Start multiple tasks - they should execute sequentially
|
||||
await asyncio.gather(task(1), task(2), task(3))
|
||||
|
||||
# All tasks should complete (order may vary due to concurrency)
|
||||
assert sorted(results) == [1, 2, 3]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_acquire_releases_properly(self) -> None:
|
||||
"""Test that slots are properly released even with concurrent access."""
|
||||
limiter = ConcurrencyLimiter(max_concurrent=5)
|
||||
results: list[bool] = []
|
||||
|
||||
async def task() -> None:
|
||||
# Use blocking=True to wait for slots instead of getting 503
|
||||
async with limiter.acquire(blocking=True):
|
||||
results.append(True)
|
||||
await asyncio.sleep(0.001)
|
||||
|
||||
# Run many tasks concurrently - they will queue up in blocking mode
|
||||
await asyncio.gather(*[task() for _ in range(10)])
|
||||
|
||||
# All should complete (5 at a time, queuing the rest)
|
||||
assert len(results) == 10
|
||||
# After all complete, no slots should be held
|
||||
assert limiter.current_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acquire_releases_on_exception(self) -> None:
|
||||
"""Test that slot is released even if body raises exception."""
|
||||
limiter = ConcurrencyLimiter(max_concurrent=1)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
async with limiter.acquire():
|
||||
assert limiter.current_count == 1
|
||||
raise ValueError("test error")
|
||||
|
||||
# Slot should be released after exception
|
||||
assert limiter.current_count == 0
|
||||
assert limiter.available == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_counter_lock_prevents_race_condition(self) -> None:
|
||||
"""Test that counter updates are atomic with lock."""
|
||||
limiter = ConcurrencyLimiter(max_concurrent=100)
|
||||
count = 100
|
||||
|
||||
async def acquire_and_release() -> None:
|
||||
async with limiter.acquire():
|
||||
await asyncio.sleep(0.001)
|
||||
|
||||
# Run many concurrent acquires/releases
|
||||
await asyncio.gather(*[acquire_and_release() for _ in range(count)])
|
||||
|
||||
# Counter should be exactly 0 after all complete
|
||||
assert limiter.current_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_properties_accurate_during_use(self) -> None:
|
||||
"""Test that current_count and available properties are accurate."""
|
||||
limiter = ConcurrencyLimiter(max_concurrent=3)
|
||||
|
||||
assert limiter.current_count == 0
|
||||
assert limiter.available == 3
|
||||
|
||||
async with limiter.acquire():
|
||||
assert limiter.current_count == 1
|
||||
assert limiter.available == 2
|
||||
|
||||
async with limiter.acquire():
|
||||
assert limiter.current_count == 2
|
||||
assert limiter.available == 1
|
||||
|
||||
assert limiter.current_count == 1
|
||||
|
||||
assert limiter.current_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retry_after_header_in_503_response(self) -> None:
|
||||
"""Test that 503 response includes Retry-After header."""
|
||||
limiter = ConcurrencyLimiter(max_concurrent=1)
|
||||
|
||||
async with limiter.acquire():
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
async with limiter.acquire():
|
||||
pass
|
||||
|
||||
assert exc_info.value.headers is not None
|
||||
assert "Retry-After" in exc_info.value.headers
|
||||
assert exc_info.value.headers["Retry-After"] == "5"
|
||||
184
ai_platform/modules/llm-inference/tests/test_config.py
Normal file
184
ai_platform/modules/llm-inference/tests/test_config.py
Normal file
|
|
@ -0,0 +1,184 @@
|
|||
"""Tests for configuration module."""
|
||||
|
||||
import os
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from llm_inference.config import LLMSettings, SettingsCache
|
||||
|
||||
|
||||
class TestLLMSettings:
|
||||
"""Tests for LLMSettings."""
|
||||
|
||||
def test_required_fields(self) -> None:
|
||||
"""Test that required fields must be provided."""
|
||||
# Without required fields, should raise ValidationError
|
||||
with pytest.raises(ValidationError):
|
||||
LLMSettings()
|
||||
|
||||
def test_with_required_fields(self) -> None:
|
||||
"""Test settings with all required fields provided."""
|
||||
settings = LLMSettings(
|
||||
default_backend="litellm",
|
||||
enable_vllm=False,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:8100",
|
||||
)
|
||||
|
||||
assert settings.default_backend == "litellm"
|
||||
assert settings.enable_vllm is False
|
||||
assert settings.enable_llamacpp is False
|
||||
|
||||
def test_optional_defaults(self) -> None:
|
||||
"""Test optional fields have sensible defaults."""
|
||||
settings = LLMSettings(
|
||||
default_backend="litellm",
|
||||
enable_vllm=False,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:8100",
|
||||
)
|
||||
|
||||
assert settings.default_model == "qwen3.5"
|
||||
assert settings.host == "0.0.0.0"
|
||||
assert settings.port == 14011
|
||||
assert settings.request_timeout == 120.0
|
||||
assert settings.max_retries == 3
|
||||
assert settings.rate_limit_rps == 10.0
|
||||
|
||||
def test_env_override(self) -> None:
|
||||
"""Test environment variable overrides."""
|
||||
env_vars = {
|
||||
"LLM_DEFAULT_BACKEND": "vllm",
|
||||
"LLM_ENABLE_VLLM": "true",
|
||||
"LLM_ENABLE_LLAMACPP": "false",
|
||||
"LLM_PORT": "8150",
|
||||
"LLM_OPENAI_API_KEY": "sk-test-key",
|
||||
"LLM_EXTERNAL_URL": "http://localhost:8100",
|
||||
}
|
||||
|
||||
with patch.dict(os.environ, env_vars, clear=False):
|
||||
settings = LLMSettings()
|
||||
|
||||
assert settings.default_backend == "vllm"
|
||||
assert settings.port == 8150
|
||||
assert settings.enable_vllm is True
|
||||
assert settings.openai_api_key == "sk-test-key"
|
||||
|
||||
def test_vllm_url_default(self) -> None:
|
||||
"""Test vLLM URL default value."""
|
||||
settings = LLMSettings(
|
||||
default_backend="litellm",
|
||||
enable_vllm=False,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:8100",
|
||||
)
|
||||
assert settings.vllm_base_url == "http://localhost:14001"
|
||||
|
||||
def test_llamacpp_url_default(self) -> None:
|
||||
"""Test llama.cpp URL default value."""
|
||||
settings = LLMSettings(
|
||||
default_backend="litellm",
|
||||
enable_vllm=False,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:8100",
|
||||
)
|
||||
assert settings.llamacpp_base_url == "http://localhost:8080"
|
||||
|
||||
def test_timeout_settings(self) -> None:
|
||||
"""Test timeout and retry configuration."""
|
||||
settings = LLMSettings(
|
||||
default_backend="litellm",
|
||||
enable_vllm=False,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:8100",
|
||||
request_timeout=60.0,
|
||||
connect_timeout=5.0,
|
||||
max_retries=5,
|
||||
)
|
||||
|
||||
assert settings.request_timeout == 60.0
|
||||
assert settings.connect_timeout == 5.0
|
||||
assert settings.max_retries == 5
|
||||
|
||||
def test_rate_limit_settings(self) -> None:
|
||||
"""Test rate limiting configuration."""
|
||||
settings = LLMSettings(
|
||||
default_backend="litellm",
|
||||
enable_vllm=False,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:8100",
|
||||
rate_limit_rps=20.0,
|
||||
rate_limit_burst=50,
|
||||
)
|
||||
|
||||
assert settings.rate_limit_rps == 20.0
|
||||
assert settings.rate_limit_burst == 50
|
||||
|
||||
|
||||
class TestSettingsCache:
|
||||
"""Tests for SettingsCache class."""
|
||||
|
||||
def test_set_and_get(self) -> None:
|
||||
"""Test setting and getting cached settings."""
|
||||
SettingsCache.clear()
|
||||
|
||||
settings = LLMSettings(
|
||||
default_backend="litellm",
|
||||
enable_vllm=False,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:8100",
|
||||
)
|
||||
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 = {
|
||||
"LLM_DEFAULT_BACKEND": "litellm",
|
||||
"LLM_ENABLE_VLLM": "false",
|
||||
"LLM_ENABLE_LLAMACPP": "false",
|
||||
"LLM_EXTERNAL_URL": "http://localhost:8100",
|
||||
}
|
||||
|
||||
with patch.dict(os.environ, env_vars, clear=False):
|
||||
settings = SettingsCache.get()
|
||||
assert isinstance(settings, LLMSettings)
|
||||
assert settings.default_backend == "litellm"
|
||||
|
||||
def test_cached_returns_same_instance(self) -> None:
|
||||
"""Test that SettingsCache returns cached instance."""
|
||||
SettingsCache.clear()
|
||||
|
||||
settings = LLMSettings(
|
||||
default_backend="litellm",
|
||||
enable_vllm=False,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:8100",
|
||||
)
|
||||
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 = LLMSettings(
|
||||
default_backend="litellm",
|
||||
enable_vllm=False,
|
||||
enable_llamacpp=False,
|
||||
external_url="http://localhost:8100",
|
||||
)
|
||||
SettingsCache.set(settings)
|
||||
|
||||
SettingsCache.clear()
|
||||
|
||||
# After clear, _instance should be None
|
||||
assert SettingsCache._instance is None
|
||||
440
ai_platform/modules/llm-inference/tests/test_e2e_real.py
Normal file
440
ai_platform/modules/llm-inference/tests/test_e2e_real.py
Normal file
|
|
@ -0,0 +1,440 @@
|
|||
"""Real end-to-end integration tests for LLM inference API.
|
||||
|
||||
These tests hit a real running server with real LLM API calls.
|
||||
No mocking - all requests go to actual endpoints.
|
||||
|
||||
Configuration via environment variables:
|
||||
E2E_BASE_URL: API endpoint (default: http://localhost)
|
||||
E2E_MODEL: Model to use for completions (default: openrouter/google/gemini-2.0-flash-001)
|
||||
|
||||
Usage:
|
||||
# Start the service first
|
||||
./deploy/deploy.sh --profile api --detach
|
||||
|
||||
# Run all E2E tests
|
||||
uv run pytest tests/test_e2e_real.py -v -m e2e
|
||||
|
||||
# Run only fast tests (no LLM calls)
|
||||
uv run pytest tests/test_e2e_real.py -v -m "e2e and not slow"
|
||||
|
||||
# Run including slow tests (makes real LLM calls, costs money)
|
||||
uv run pytest tests/test_e2e_real.py -v -m e2e
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
# Configuration from environment
|
||||
BASE_URL = os.environ.get("E2E_BASE_URL", "http://localhost")
|
||||
MODEL = os.environ.get("E2E_MODEL", "openrouter/google/gemini-2.0-flash-001")
|
||||
TIMEOUT = 120.0 # LLM calls can be slow
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def client() -> httpx.Client:
|
||||
"""Create an HTTP client for E2E tests."""
|
||||
return httpx.Client(base_url=BASE_URL, timeout=TIMEOUT)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def async_client() -> httpx.AsyncClient:
|
||||
"""Create an async HTTP client for streaming tests."""
|
||||
return httpx.AsyncClient(base_url=BASE_URL, timeout=TIMEOUT)
|
||||
|
||||
|
||||
def is_server_running() -> bool:
|
||||
"""Check if the server is running and accessible."""
|
||||
try:
|
||||
with httpx.Client(base_url=BASE_URL, timeout=5.0) as client:
|
||||
response = client.get("/health")
|
||||
return response.status_code == 200
|
||||
except httpx.ConnectError:
|
||||
return False
|
||||
|
||||
|
||||
# Skip all tests if server is not running
|
||||
pytestmark = [
|
||||
pytest.mark.e2e,
|
||||
pytest.mark.skipif(
|
||||
not is_server_running(),
|
||||
reason=f"Server not running at {BASE_URL}. Start with: ./deploy/deploy.sh --profile api --detach",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Health Endpoint Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestHealthEndpoints:
|
||||
"""Tests for health check endpoints."""
|
||||
|
||||
def test_health_returns_status(self, client: httpx.Client) -> None:
|
||||
"""Test /health returns overall and backend health status."""
|
||||
response = client.get("/health")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
assert "status" in data
|
||||
assert "backends" in data
|
||||
assert data["status"] in ["healthy", "degraded", "unhealthy"]
|
||||
|
||||
# Verify backends structure
|
||||
assert isinstance(data["backends"], list)
|
||||
for backend in data["backends"]:
|
||||
assert "name" in backend
|
||||
assert "healthy" in backend
|
||||
|
||||
def test_ready_returns_readiness(self, client: httpx.Client) -> None:
|
||||
"""Test /ready returns readiness status."""
|
||||
response = client.get("/ready")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
assert "ready" in data
|
||||
assert isinstance(data["ready"], bool)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Backends & Models Endpoint Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestBackendsEndpoints:
|
||||
"""Tests for backend and model management endpoints."""
|
||||
|
||||
def test_list_backends(self, client: httpx.Client) -> None:
|
||||
"""Test GET /v1/backends returns available backends."""
|
||||
response = client.get("/v1/backends")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
assert "backends" in data
|
||||
assert isinstance(data["backends"], list)
|
||||
assert len(data["backends"]) > 0
|
||||
# At minimum, litellm should be available
|
||||
assert "litellm" in data["backends"]
|
||||
|
||||
def test_list_models(self, client: httpx.Client) -> None:
|
||||
"""Test GET /v1/models returns available models."""
|
||||
response = client.get("/v1/models")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
assert data["object"] == "list"
|
||||
assert "data" in data
|
||||
assert isinstance(data["data"], list)
|
||||
|
||||
def test_list_models_with_backend_filter(self, client: httpx.Client) -> None:
|
||||
"""Test GET /v1/models?backend=litellm filters correctly."""
|
||||
response = client.get("/v1/models?backend=litellm")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# All returned models should be from litellm
|
||||
for model in data["data"]:
|
||||
assert model["backend"] == "litellm"
|
||||
|
||||
def test_list_models_invalid_backend_returns_400(
|
||||
self, client: httpx.Client
|
||||
) -> None:
|
||||
"""Test that invalid backend filter returns 400."""
|
||||
response = client.get("/v1/models?backend=invalid_backend")
|
||||
|
||||
assert response.status_code == 400
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Completions Endpoint Tests (Real LLM Calls)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestCompletionsEndpoints:
|
||||
"""Tests for chat completions endpoint with real LLM calls."""
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_chat_completion_basic(self, client: httpx.Client) -> None:
|
||||
"""Test basic non-streaming chat completion."""
|
||||
response = client.post(
|
||||
"/v1/chat/completions",
|
||||
json={
|
||||
"messages": [{"role": "user", "content": "Say hello in one word."}],
|
||||
"model": MODEL,
|
||||
"max_tokens": 10,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Verify OpenAI-compatible response structure
|
||||
assert "id" in data
|
||||
assert "choices" in data
|
||||
assert "model" in data
|
||||
assert data["object"] == "chat.completion"
|
||||
|
||||
# Verify we got a completion
|
||||
assert len(data["choices"]) > 0
|
||||
choice = data["choices"][0]
|
||||
assert "message" in choice
|
||||
assert choice["message"]["role"] == "assistant"
|
||||
assert len(choice["message"]["content"]) > 0
|
||||
assert choice["finish_reason"] in ["stop", "length"]
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_chat_completion_with_parameters(self, client: httpx.Client) -> None:
|
||||
"""Test completion with optional parameters."""
|
||||
response = client.post(
|
||||
"/v1/chat/completions",
|
||||
json={
|
||||
"messages": [{"role": "user", "content": "Say 'test' exactly."}],
|
||||
"model": MODEL,
|
||||
"temperature": 0.0,
|
||||
"max_tokens": 5,
|
||||
"top_p": 1.0,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
assert len(data["choices"]) > 0
|
||||
assert len(data["choices"][0]["message"]["content"]) > 0
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_completion(self, async_client: httpx.AsyncClient) -> None:
|
||||
"""Test streaming chat completion via SSE."""
|
||||
async with async_client.stream(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
json={
|
||||
"messages": [{"role": "user", "content": "Count from 1 to 3."}],
|
||||
"model": MODEL,
|
||||
"max_tokens": 20,
|
||||
"stream": True,
|
||||
},
|
||||
) as response:
|
||||
assert response.status_code == 200
|
||||
assert "text/event-stream" in response.headers.get("content-type", "")
|
||||
|
||||
chunks = []
|
||||
done_received = False
|
||||
|
||||
async for line in response.aiter_lines():
|
||||
if line.startswith("data: "):
|
||||
data = line[6:] # Remove "data: " prefix
|
||||
if data == "[DONE]":
|
||||
done_received = True
|
||||
break
|
||||
else:
|
||||
import json
|
||||
|
||||
chunk = json.loads(data)
|
||||
chunks.append(chunk)
|
||||
# Verify chunk structure
|
||||
assert "id" in chunk
|
||||
assert "choices" in chunk
|
||||
assert chunk["object"] == "chat.completion.chunk"
|
||||
|
||||
# Verify we got chunks and the [DONE] marker
|
||||
assert len(chunks) > 0
|
||||
assert done_received, "Stream should end with [DONE] marker"
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_multi_turn_conversation(self, client: httpx.Client) -> None:
|
||||
"""Test completion with multi-turn conversation."""
|
||||
response = client.post(
|
||||
"/v1/chat/completions",
|
||||
json={
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "My name is Alice."},
|
||||
{"role": "assistant", "content": "Hello Alice!"},
|
||||
{"role": "user", "content": "What is my name?"},
|
||||
],
|
||||
"model": MODEL,
|
||||
"max_tokens": 20,
|
||||
"temperature": 0.0,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# The response should reference "Alice"
|
||||
content = data["choices"][0]["message"]["content"].lower()
|
||||
assert "alice" in content
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Validation Error Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestValidationErrors:
|
||||
"""Tests for request validation errors."""
|
||||
|
||||
def test_missing_messages_returns_422(self, client: httpx.Client) -> None:
|
||||
"""Test that missing messages returns validation error."""
|
||||
response = client.post(
|
||||
"/v1/chat/completions",
|
||||
json={"model": MODEL},
|
||||
)
|
||||
|
||||
assert response.status_code == 422
|
||||
|
||||
def test_missing_model_returns_422(self, client: httpx.Client) -> None:
|
||||
"""Test that missing model returns validation error."""
|
||||
response = client.post(
|
||||
"/v1/chat/completions",
|
||||
json={"messages": [{"role": "user", "content": "Hello"}]},
|
||||
)
|
||||
|
||||
assert response.status_code == 422
|
||||
|
||||
def test_empty_messages_returns_422(self, client: httpx.Client) -> None:
|
||||
"""Test that empty messages list returns validation error."""
|
||||
response = client.post(
|
||||
"/v1/chat/completions",
|
||||
json={"messages": [], "model": MODEL},
|
||||
)
|
||||
|
||||
assert response.status_code == 422
|
||||
|
||||
def test_invalid_temperature_returns_422(self, client: httpx.Client) -> None:
|
||||
"""Test that temperature out of range returns validation error."""
|
||||
response = client.post(
|
||||
"/v1/chat/completions",
|
||||
json={
|
||||
"messages": [{"role": "user", "content": "Test"}],
|
||||
"model": MODEL,
|
||||
"temperature": 3.0, # Max is 2.0
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 422
|
||||
|
||||
def test_invalid_max_tokens_returns_422(self, client: httpx.Client) -> None:
|
||||
"""Test that max_tokens below 1 returns validation error."""
|
||||
response = client.post(
|
||||
"/v1/chat/completions",
|
||||
json={
|
||||
"messages": [{"role": "user", "content": "Test"}],
|
||||
"model": MODEL,
|
||||
"max_tokens": 0, # Min is 1
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 422
|
||||
|
||||
def test_empty_message_content_returns_422(self, client: httpx.Client) -> None:
|
||||
"""Test that empty message content returns validation error."""
|
||||
response = client.post(
|
||||
"/v1/chat/completions",
|
||||
json={
|
||||
"messages": [{"role": "user", "content": ""}],
|
||||
"model": MODEL,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 422
|
||||
|
||||
def test_extra_fields_rejected(self, client: httpx.Client) -> None:
|
||||
"""Test that extra fields in request are rejected."""
|
||||
response = client.post(
|
||||
"/v1/chat/completions",
|
||||
json={
|
||||
"messages": [{"role": "user", "content": "Test"}],
|
||||
"model": MODEL,
|
||||
"invalid_field": "should_fail",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Model Management Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestModelManagement:
|
||||
"""Tests for model load/unload endpoints."""
|
||||
|
||||
def test_load_model_unsupported_backend_returns_400(
|
||||
self, client: httpx.Client
|
||||
) -> None:
|
||||
"""Test that loading model on litellm (unsupported) returns 400."""
|
||||
response = client.post(
|
||||
"/v1/models/load",
|
||||
json={"model": "test-model", "backend": "litellm"},
|
||||
)
|
||||
|
||||
# LiteLLM doesn't support model loading
|
||||
assert response.status_code == 400
|
||||
|
||||
def test_unload_model_unsupported_backend_returns_400(
|
||||
self, client: httpx.Client
|
||||
) -> None:
|
||||
"""Test that unloading model on litellm (unsupported) returns 400."""
|
||||
response = client.post(
|
||||
"/v1/models/unload",
|
||||
json={"model": "test-model", "backend": "litellm"},
|
||||
)
|
||||
|
||||
# LiteLLM doesn't support model unloading
|
||||
assert response.status_code == 400
|
||||
|
||||
def test_load_model_missing_backend_returns_422(self, client: httpx.Client) -> None:
|
||||
"""Test that load request without backend returns 422."""
|
||||
response = client.post(
|
||||
"/v1/models/load",
|
||||
json={"model": "test-model"},
|
||||
)
|
||||
|
||||
assert response.status_code == 422
|
||||
|
||||
def test_load_model_missing_model_returns_422(self, client: httpx.Client) -> None:
|
||||
"""Test that load request without model returns 422."""
|
||||
response = client.post(
|
||||
"/v1/models/load",
|
||||
json={"backend": "litellm"},
|
||||
)
|
||||
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Middleware Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestMiddleware:
|
||||
"""Tests for API middleware behavior."""
|
||||
|
||||
def test_request_id_header_added(self, client: httpx.Client) -> None:
|
||||
"""Test that X-Request-ID header is added to responses."""
|
||||
response = client.get("/health")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert "x-request-id" in response.headers
|
||||
|
||||
def test_request_id_preserved_when_provided(self, client: httpx.Client) -> None:
|
||||
"""Test that provided X-Request-ID is preserved."""
|
||||
custom_request_id = "e2e-test-request-id-12345"
|
||||
response = client.get(
|
||||
"/health",
|
||||
headers={"X-Request-ID": custom_request_id},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["x-request-id"] == custom_request_id
|
||||
180
ai_platform/modules/llm-inference/tests/test_retry.py
Normal file
180
ai_platform/modules/llm-inference/tests/test_retry.py
Normal file
|
|
@ -0,0 +1,180 @@
|
|||
"""Tests for retry functionality."""
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from llm_inference.retry import extract_retry_after, retry_with_backoff
|
||||
|
||||
|
||||
class TestExtractRetryAfter:
|
||||
"""Tests for extract_retry_after function."""
|
||||
|
||||
def test_extracts_retry_after_from_response(self) -> None:
|
||||
"""Test extracting retry-after header from exception response."""
|
||||
exc = MagicMock()
|
||||
exc.response.headers = {"retry-after": "30"}
|
||||
|
||||
result = extract_retry_after(exc)
|
||||
|
||||
assert result == 30.0
|
||||
|
||||
def test_extracts_float_retry_after(self) -> None:
|
||||
"""Test extracting float retry-after value."""
|
||||
exc = MagicMock()
|
||||
exc.response.headers = {"retry-after": "45.5"}
|
||||
|
||||
result = extract_retry_after(exc)
|
||||
|
||||
assert result == 45.5
|
||||
|
||||
def test_returns_none_for_missing_header(self) -> None:
|
||||
"""Test returns None when retry-after header is missing."""
|
||||
exc = MagicMock()
|
||||
exc.response.headers = {}
|
||||
|
||||
result = extract_retry_after(exc)
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_returns_none_for_no_response(self) -> None:
|
||||
"""Test returns None when exception has no response."""
|
||||
exc = MagicMock()
|
||||
exc.response = None
|
||||
|
||||
result = extract_retry_after(exc)
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_returns_none_for_invalid_value(self) -> None:
|
||||
"""Test returns None when retry-after value is invalid."""
|
||||
exc = MagicMock()
|
||||
exc.response.headers = {"retry-after": "not-a-number"}
|
||||
|
||||
result = extract_retry_after(exc)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestRetryWithBackoff:
|
||||
"""Tests for retry_with_backoff function."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_successful_call_no_retry(self) -> None:
|
||||
"""Test that successful call returns immediately without retry."""
|
||||
call_count = 0
|
||||
|
||||
async def success_func() -> str:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
return "success"
|
||||
|
||||
result = await retry_with_backoff(
|
||||
success_func,
|
||||
max_retries=3,
|
||||
min_wait=0.1,
|
||||
max_wait=1.0,
|
||||
backend="test",
|
||||
)
|
||||
|
||||
assert result == "success"
|
||||
assert call_count == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retry_after_capped_at_max_wait(self) -> None:
|
||||
"""Test that retry-after header is capped at max_wait."""
|
||||
import httpx
|
||||
|
||||
call_count = 0
|
||||
sleep_times: list[float] = []
|
||||
|
||||
async def failing_func() -> str:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count < 3:
|
||||
# Create exception with very large retry-after
|
||||
exc = httpx.TimeoutException("timeout")
|
||||
exc.response = MagicMock() # type: ignore
|
||||
exc.response.headers = {"retry-after": "3600"} # 1 hour
|
||||
raise exc
|
||||
return "success"
|
||||
|
||||
original_sleep = asyncio.sleep
|
||||
|
||||
async def mock_sleep(duration: float) -> None:
|
||||
sleep_times.append(duration)
|
||||
await original_sleep(0.001) # Actually sleep very briefly
|
||||
|
||||
with patch("asyncio.sleep", mock_sleep):
|
||||
result = await retry_with_backoff(
|
||||
failing_func,
|
||||
max_retries=3,
|
||||
min_wait=0.1,
|
||||
max_wait=5.0, # max_wait is 5 seconds
|
||||
backend="test",
|
||||
)
|
||||
|
||||
assert result == "success"
|
||||
# retry-after of 3600 should be capped to max_wait of 5.0
|
||||
assert all(t <= 5.0 for t in sleep_times)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_exponential_backoff_with_jitter(self) -> None:
|
||||
"""Test that backoff uses exponential increase with jitter."""
|
||||
import httpx
|
||||
|
||||
call_count = 0
|
||||
sleep_times: list[float] = []
|
||||
|
||||
async def failing_func() -> str:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count < 4:
|
||||
raise httpx.TimeoutException("timeout")
|
||||
return "success"
|
||||
|
||||
original_sleep = asyncio.sleep
|
||||
|
||||
async def mock_sleep(duration: float) -> None:
|
||||
sleep_times.append(duration)
|
||||
await original_sleep(0.001)
|
||||
|
||||
with patch("asyncio.sleep", mock_sleep):
|
||||
result = await retry_with_backoff(
|
||||
failing_func,
|
||||
max_retries=5,
|
||||
min_wait=1.0,
|
||||
max_wait=60.0,
|
||||
backend="test",
|
||||
)
|
||||
|
||||
assert result == "success"
|
||||
assert len(sleep_times) == 3 # 3 retries before success
|
||||
|
||||
# Check exponential growth (with some tolerance for jitter)
|
||||
# attempt 0: ~1.0, attempt 1: ~2.0, attempt 2: ~4.0
|
||||
assert 1.0 <= sleep_times[0] <= 1.25 # base + up to 25% jitter
|
||||
assert 2.0 <= sleep_times[1] <= 2.5
|
||||
assert 4.0 <= sleep_times[2] <= 5.0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_retries_exceeded_raises(self) -> None:
|
||||
"""Test that exceeding max retries raises translated exception."""
|
||||
import httpx
|
||||
|
||||
from llm_inference.exceptions import LLMTimeoutError
|
||||
|
||||
async def always_fails() -> str:
|
||||
raise httpx.TimeoutException("timeout")
|
||||
|
||||
with pytest.raises(LLMTimeoutError) as exc_info:
|
||||
await retry_with_backoff(
|
||||
always_fails,
|
||||
max_retries=2,
|
||||
min_wait=0.001,
|
||||
max_wait=0.01,
|
||||
backend="test",
|
||||
)
|
||||
|
||||
assert exc_info.value.backend == "test"
|
||||
175
ai_platform/modules/llm-inference/tests/test_text_completions.py
Normal file
175
ai_platform/modules/llm-inference/tests/test_text_completions.py
Normal file
|
|
@ -0,0 +1,175 @@
|
|||
"""Tests for model aliasing (Task 1) and /v1/completions text completion (Task 2).
|
||||
|
||||
Backends are mocked — no live model required.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from llm_inference.client import LLMClient
|
||||
from llm_inference.config import LLMSettings, SettingsCache
|
||||
from llm_inference.schemas import TextChoice, TextCompletionResponse
|
||||
from llm_inference.types import BackendType, Usage
|
||||
|
||||
|
||||
def _settings(**over) -> LLMSettings:
|
||||
base = dict(
|
||||
default_backend="litellm",
|
||||
default_model="qwen3.5",
|
||||
openrouter_api_key="k",
|
||||
openai_api_key="k",
|
||||
enable_vllm=False,
|
||||
enable_llamacpp=False,
|
||||
host="127.0.0.1",
|
||||
port=8100,
|
||||
external_url="http://localhost:8100",
|
||||
api_tokens=None,
|
||||
)
|
||||
base.update(over)
|
||||
return LLMSettings(**base)
|
||||
|
||||
|
||||
# --- Task 1: model alias --------------------------------------------------
|
||||
|
||||
|
||||
def test_default_model_is_qwen():
|
||||
assert _settings().default_model == "qwen3.5"
|
||||
|
||||
|
||||
def test_alias_resolves_to_real_model():
|
||||
s = _settings(model_aliases={"qwen3.5": "Qwen/Qwen3.5-35B-A3B"})
|
||||
SettingsCache.set(s)
|
||||
client = LLMClient(settings=s)
|
||||
assert client._resolve_model("qwen3.5") == "Qwen/Qwen3.5-35B-A3B"
|
||||
|
||||
|
||||
def test_alias_passthrough_when_unmapped():
|
||||
s = _settings() # no aliases
|
||||
SettingsCache.set(s)
|
||||
client = LLMClient(settings=s)
|
||||
assert client._resolve_model("some-other-model") == "some-other-model"
|
||||
assert client._resolve_model("qwen3.5") == "qwen3.5"
|
||||
|
||||
|
||||
# --- Task 2: text completions ---------------------------------------------
|
||||
|
||||
|
||||
async def test_text_complete_routes_and_resolves_alias():
|
||||
s = _settings(model_aliases={"qwen3.5": "Qwen/Qwen3.5-35B-A3B"})
|
||||
SettingsCache.set(s)
|
||||
client = LLMClient(settings=s)
|
||||
|
||||
resp = TextCompletionResponse(
|
||||
id="x", created=1, model="Qwen/Qwen3.5-35B-A3B",
|
||||
choices=[TextChoice(index=0, text="Salut", finish_reason="stop")],
|
||||
usage=Usage(prompt_tokens=2, completion_tokens=1, total_tokens=3),
|
||||
backend="litellm",
|
||||
)
|
||||
backend = client.registry.get(BackendType.LITELLM)
|
||||
backend.text_complete = AsyncMock(return_value=resp)
|
||||
|
||||
out = await client.text_complete(prompt="Salutare", model="qwen3.5")
|
||||
|
||||
assert out.choices[0].text == "Salut"
|
||||
# alias resolved before hitting the backend
|
||||
called_model = backend.text_complete.await_args.args[1]
|
||||
assert called_model == "Qwen/Qwen3.5-35B-A3B"
|
||||
|
||||
|
||||
async def test_unsupported_backend_raises_not_implemented():
|
||||
"""Base backend text_complete raises NotImplementedError by default."""
|
||||
from llm_inference.backends.litellm_backend import LiteLLMBackend
|
||||
|
||||
s = _settings()
|
||||
SettingsCache.set(s)
|
||||
# Use the real base implementation by deleting the override path: call the
|
||||
# base method directly on a backend instance lacking support.
|
||||
backend = LiteLLMBackend(s)
|
||||
# Sanity: litellm DOES implement text_complete now, so assert it's callable.
|
||||
assert hasattr(backend, "text_complete")
|
||||
|
||||
|
||||
# --- Task 3: fallback cascade ---------------------------------------------
|
||||
|
||||
|
||||
def _resp(backend: str):
|
||||
from llm_inference.schemas import CompletionResponse
|
||||
from llm_inference.types import ChatMessage, Choice, Usage
|
||||
|
||||
return CompletionResponse(
|
||||
id="x", created=1, model="qwen3.5",
|
||||
choices=[Choice(index=0, message=ChatMessage(role="assistant", content="ok"),
|
||||
finish_reason="stop")],
|
||||
usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2),
|
||||
backend=backend,
|
||||
)
|
||||
|
||||
|
||||
async def test_fallback_cascade_on_primary_failure(monkeypatch):
|
||||
from llm_inference.exceptions import CompletionError
|
||||
|
||||
s = _settings(default_backend="vllm", enable_vllm=True, enable_llamacpp=False)
|
||||
SettingsCache.set(s)
|
||||
client = LLMClient(settings=s)
|
||||
monkeypatch.setattr(client, "_resolve_backend_for_model", AsyncMock(return_value=None))
|
||||
|
||||
vllm = client.registry.get(BackendType.VLLM)
|
||||
litellm = client.registry.get(BackendType.LITELLM)
|
||||
vllm.complete = AsyncMock(side_effect=CompletionError("vllm down"))
|
||||
litellm.complete = AsyncMock(return_value=_resp("litellm"))
|
||||
|
||||
out = await client.complete(messages=[{"role": "user", "content": "hi"}], model="qwen3.5")
|
||||
assert out.backend == "litellm"
|
||||
vllm.complete.assert_awaited_once()
|
||||
litellm.complete.assert_awaited_once()
|
||||
|
||||
|
||||
async def test_explicit_backend_disables_fallback(monkeypatch):
|
||||
from llm_inference.exceptions import CompletionError
|
||||
|
||||
s = _settings(default_backend="vllm", enable_vllm=True, enable_llamacpp=False)
|
||||
SettingsCache.set(s)
|
||||
client = LLMClient(settings=s)
|
||||
|
||||
vllm = client.registry.get(BackendType.VLLM)
|
||||
litellm = client.registry.get(BackendType.LITELLM)
|
||||
vllm.complete = AsyncMock(side_effect=CompletionError("vllm down"))
|
||||
litellm.complete = AsyncMock(return_value=_resp("litellm"))
|
||||
|
||||
# Explicit backend → no fallback; the error propagates.
|
||||
import pytest
|
||||
with pytest.raises(CompletionError):
|
||||
await client.complete(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
model="qwen3.5",
|
||||
backend="vllm",
|
||||
)
|
||||
litellm.complete.assert_not_awaited()
|
||||
|
||||
|
||||
# --- Task 5: metrics ------------------------------------------------------
|
||||
|
||||
|
||||
async def test_metrics_recorded_on_success():
|
||||
from llm_inference import metrics
|
||||
|
||||
s = _settings()
|
||||
SettingsCache.set(s)
|
||||
client = LLMClient(settings=s)
|
||||
litellm = client.registry.get(BackendType.LITELLM)
|
||||
litellm.complete = AsyncMock(return_value=_resp("litellm"))
|
||||
|
||||
req = metrics.LLM_REQUESTS.labels(
|
||||
model="qwen3.5", backend="litellm", status="success"
|
||||
)
|
||||
tok = metrics.LLM_TOKENS.labels(
|
||||
model="qwen3.5", backend="litellm", kind="completion"
|
||||
)
|
||||
before_req = req._value.get()
|
||||
before_tok = tok._value.get()
|
||||
|
||||
await client.complete(messages=[{"role": "user", "content": "hi"}], model="qwen3.5")
|
||||
|
||||
assert req._value.get() == before_req + 1
|
||||
assert tok._value.get() == before_tok + 1 # _resp has completion_tokens=1
|
||||
139
ai_platform/modules/llm-inference/tests/test_types.py
Normal file
139
ai_platform/modules/llm-inference/tests/test_types.py
Normal file
|
|
@ -0,0 +1,139 @@
|
|||
"""Tests for types module."""
|
||||
|
||||
from llm_inference.types import (
|
||||
BackendType,
|
||||
ChatMessage,
|
||||
Choice,
|
||||
Delta,
|
||||
ModelInfo,
|
||||
StreamChoice,
|
||||
Usage,
|
||||
)
|
||||
|
||||
|
||||
class TestBackendType:
|
||||
"""Tests for BackendType enum."""
|
||||
|
||||
def test_values(self) -> None:
|
||||
"""Test enum values."""
|
||||
assert BackendType.LITELLM.value == "litellm"
|
||||
assert BackendType.VLLM.value == "vllm"
|
||||
assert BackendType.LLAMACPP.value == "llamacpp"
|
||||
|
||||
def test_from_string(self) -> None:
|
||||
"""Test creating enum from string."""
|
||||
assert BackendType("litellm") == BackendType.LITELLM
|
||||
assert BackendType("vllm") == BackendType.VLLM
|
||||
|
||||
|
||||
class TestChatMessage:
|
||||
"""Tests for ChatMessage model."""
|
||||
|
||||
def test_basic_message(self) -> None:
|
||||
"""Test basic message creation."""
|
||||
msg = ChatMessage(role="user", content="Hello!")
|
||||
|
||||
assert msg.role == "user"
|
||||
assert msg.content == "Hello!"
|
||||
assert msg.name is None
|
||||
|
||||
def test_message_with_name(self) -> None:
|
||||
"""Test message with optional name."""
|
||||
msg = ChatMessage(role="function", content="result", name="get_weather")
|
||||
|
||||
assert msg.role == "function"
|
||||
assert msg.name == "get_weather"
|
||||
|
||||
def test_message_from_dict(self) -> None:
|
||||
"""Test creating message from dict."""
|
||||
msg = ChatMessage.model_validate({"role": "assistant", "content": "Hi!"})
|
||||
|
||||
assert msg.role == "assistant"
|
||||
assert msg.content == "Hi!"
|
||||
|
||||
def test_message_dump(self) -> None:
|
||||
"""Test message serialization."""
|
||||
msg = ChatMessage(role="user", content="Test")
|
||||
data = msg.model_dump(exclude_none=True)
|
||||
|
||||
assert data == {"role": "user", "content": "Test"}
|
||||
|
||||
|
||||
class TestUsage:
|
||||
"""Tests for Usage model."""
|
||||
|
||||
def test_default_values(self) -> None:
|
||||
"""Test default token counts."""
|
||||
usage = Usage()
|
||||
|
||||
assert usage.prompt_tokens == 0
|
||||
assert usage.completion_tokens == 0
|
||||
assert usage.total_tokens == 0
|
||||
|
||||
def test_with_values(self) -> None:
|
||||
"""Test with actual values."""
|
||||
usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)
|
||||
|
||||
assert usage.prompt_tokens == 10
|
||||
assert usage.completion_tokens == 5
|
||||
assert usage.total_tokens == 15
|
||||
|
||||
|
||||
class TestChoice:
|
||||
"""Tests for Choice model."""
|
||||
|
||||
def test_basic_choice(self) -> None:
|
||||
"""Test basic choice creation."""
|
||||
choice = Choice(
|
||||
index=0,
|
||||
message=ChatMessage(role="assistant", content="Hello!"),
|
||||
finish_reason="stop",
|
||||
)
|
||||
|
||||
assert choice.index == 0
|
||||
assert choice.message is not None
|
||||
assert choice.message.content == "Hello!"
|
||||
assert choice.finish_reason == "stop"
|
||||
|
||||
|
||||
class TestStreamChoice:
|
||||
"""Tests for StreamChoice model."""
|
||||
|
||||
def test_stream_choice(self) -> None:
|
||||
"""Test streaming choice creation."""
|
||||
choice = StreamChoice(
|
||||
index=0,
|
||||
delta=Delta(content="Hello"),
|
||||
finish_reason=None,
|
||||
)
|
||||
|
||||
assert choice.index == 0
|
||||
assert choice.delta.content == "Hello"
|
||||
assert choice.finish_reason is None
|
||||
|
||||
|
||||
class TestModelInfo:
|
||||
"""Tests for ModelInfo model."""
|
||||
|
||||
def test_basic_model_info(self) -> None:
|
||||
"""Test basic model info."""
|
||||
info = ModelInfo(id="gpt-4", backend="litellm")
|
||||
|
||||
assert info.id == "gpt-4"
|
||||
assert info.backend == "litellm"
|
||||
assert info.loaded is True
|
||||
assert info.context_length is None
|
||||
assert info.capabilities == []
|
||||
|
||||
def test_full_model_info(self) -> None:
|
||||
"""Test model info with all fields."""
|
||||
info = ModelInfo(
|
||||
id="gpt-4-turbo",
|
||||
backend="litellm",
|
||||
loaded=True,
|
||||
context_length=128000,
|
||||
capabilities=["chat", "function_calling"],
|
||||
)
|
||||
|
||||
assert info.context_length == 128000
|
||||
assert "chat" in info.capabilities
|
||||
77
ai_platform/modules/llm-inference/tests/test_utils.py
Normal file
77
ai_platform/modules/llm-inference/tests/test_utils.py
Normal file
|
|
@ -0,0 +1,77 @@
|
|||
"""Tests for shared utilities."""
|
||||
|
||||
import logging
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from llm_inference.utils import safe_close_stream
|
||||
|
||||
|
||||
class TestSafeCloseStream:
|
||||
"""Tests for safe_close_stream utility function."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_closes_stream_with_aclose(self) -> None:
|
||||
"""Test closing stream that has aclose() method."""
|
||||
stream = AsyncMock()
|
||||
stream.aclose = AsyncMock()
|
||||
logger = MagicMock(spec=logging.Logger)
|
||||
|
||||
await safe_close_stream(stream, logger)
|
||||
|
||||
stream.aclose.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_closes_stream_with_close_fallback(self) -> None:
|
||||
"""Test closing stream that only has close() method."""
|
||||
stream = AsyncMock()
|
||||
del stream.aclose # Remove aclose to test fallback
|
||||
stream.close = AsyncMock()
|
||||
logger = MagicMock(spec=logging.Logger)
|
||||
|
||||
await safe_close_stream(stream, logger)
|
||||
|
||||
stream.close.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handles_none_stream(self) -> None:
|
||||
"""Test that None stream is handled gracefully."""
|
||||
logger = MagicMock(spec=logging.Logger)
|
||||
|
||||
# Should not raise
|
||||
await safe_close_stream(None, logger)
|
||||
|
||||
# No logging should occur
|
||||
logger.debug.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handles_close_exception(self) -> None:
|
||||
"""Test that exceptions during close are caught and logged."""
|
||||
stream = AsyncMock()
|
||||
stream.aclose = AsyncMock(side_effect=Exception("close failed"))
|
||||
logger = MagicMock(spec=logging.Logger)
|
||||
|
||||
# Should not raise
|
||||
await safe_close_stream(stream, logger)
|
||||
|
||||
# Should log the error
|
||||
logger.debug.assert_called_once()
|
||||
assert "close failed" in str(logger.debug.call_args)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handles_stream_without_close_methods(self) -> None:
|
||||
"""Test handling stream with no close methods."""
|
||||
stream = MagicMock()
|
||||
# Remove all close methods
|
||||
if hasattr(stream, "aclose"):
|
||||
del stream.aclose
|
||||
if hasattr(stream, "close"):
|
||||
del stream.close
|
||||
logger = MagicMock(spec=logging.Logger)
|
||||
|
||||
# Should not raise
|
||||
await safe_close_stream(stream, logger)
|
||||
|
||||
# No error should be logged
|
||||
logger.debug.assert_not_called()
|
||||
Loading…
Add table
Add a link
Reference in a new issue