ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add custom auth header support and increase default prompt size to 100k chars (#19436)

Browse Source
This commit is contained in:
Alexsander Hamir 2026-01-20 13:25:12 -08:00 committed by GitHub
parent e142474e0b
commit 7f81dea8b3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 151 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

56
scripts/health_check/health_check_client.py
View File

@ -20,6 +20,17 @@ from typing import Dict, List, Optional, Tuple
import httpx
import yaml
# Default prompt for health checks - exactly 100k characters
# Generate a repeating pattern to reach exactly 100,000 characters
_base_text = "This is a health check test prompt for LiteLLM proxy. "
_repeat_count = (100000 // len(_base_text)) + 1
_DEFAULT_COMPLETION_PROMPT = (_base_text * _repeat_count)[:100000]
# Default embedding text - also exactly 100k characters
_embedding_base_text = "This is a test for vectorization. "
_embedding_repeat_count = (100000 // len(_embedding_base_text)) + 1
_DEFAULT_EMBEDDING_TEXT = (_embedding_base_text * _embedding_repeat_count)[:100000]
class LiteLLMHealthCheckClient:
"""Client for health checking LiteLLM proxy models."""
@ -29,8 +40,9 @@ class LiteLLMHealthCheckClient:
base_url: str,
api_key: str,
timeout: int = 120, # Match Go implementation's 120s timeout
completion_prompt: str = "Say this is a test", # Match Go implementation
embedding_text: str = "This is a test for vectorization.", # Match Go implementation
completion_prompt: str = _DEFAULT_COMPLETION_PROMPT, # Default ~100k chars
embedding_text: str = _DEFAULT_EMBEDDING_TEXT, # Default ~100k chars
custom_auth_header: Optional[str] = None,
):
"""
Initialize the health check client.
@ -41,16 +53,34 @@ class LiteLLMHealthCheckClient:
timeout: Request timeout in seconds (default: 120, matching Go implementation)
completion_prompt: Test prompt for chat/completion models
embedding_text: Test text for embedding models
custom_auth_header: Optional custom header name for authentication (e.g., "x-ifood-requester-service").
If provided, uses this header instead of standard "Authorization" header.
"""
self.base_url = base_url.rstrip("/")
self.api_key = api_key
self.timeout = timeout
self.completion_prompt = completion_prompt
self.embedding_text = embedding_text
self.headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
}
# Debug: Print prompt/text lengths
print(f"DEBUG: Completion prompt length: {len(self.completion_prompt)} characters", file=sys.stderr)
print(f"DEBUG: Embedding text length: {len(self.embedding_text)} characters", file=sys.stderr)
# Support custom auth header for proxies with custom authentication
# Handle both None and empty string
if custom_auth_header and custom_auth_header.strip():
custom_auth_header = custom_auth_header.strip()
self.headers = {
custom_auth_header: f"Bearer {api_key}",
"Content-Type": "application/json",
}
print(f"Using custom auth header: {custom_auth_header}", file=sys.stderr)
else:
self.headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
}
print("Using standard Authorization header", file=sys.stderr)
def load_models_from_yaml(self, yaml_path: str) -> List[Dict]:
"""
@ -182,6 +212,8 @@ class LiteLLMHealthCheckClient:
if is_embedding:
# Test embedding endpoint (matching Go implementation)
embedding_text_length = len(self.embedding_text)
print(f"DEBUG: Sending embedding text of length {embedding_text_length} chars to model {model_id}", file=sys.stderr)
embedding_response = await client.post(
f"{self.base_url}/v1/embeddings",
headers=self.headers,
@ -202,6 +234,8 @@ class LiteLLMHealthCheckClient:
result["dimensions"] = dimensions
else:
# Test chat completion endpoint (matching Go implementation)
prompt_length = len(self.completion_prompt)
print(f"DEBUG: Sending prompt of length {prompt_length} chars to model {model_id}", file=sys.stderr)
completion_response = await client.post(
f"{self.base_url}/v1/chat/completions",
headers=self.headers,
@ -358,6 +392,11 @@ async def main():
base_url = os.environ.get("LITELLM_BASE_URL", "http://localhost:4000")
api_key = os.environ.get("LITELLM_API_KEY", "sk-1234")
yaml_path = os.environ.get("LITELLM_MODELS_YAML")
custom_auth_header = os.environ.get("LITELLM_CUSTOM_AUTH_HEADER") # e.g., "x-ifood-requester-service"
# Debug: Print custom auth header value if set
if custom_auth_header:
print(f"Custom auth header from env: '{custom_auth_header}'", file=sys.stderr)
if not base_url:
print("Error: LITELLM_BASE_URL environment variable not set", file=sys.stderr)
@ -369,10 +408,10 @@ async def main():
timeout = int(os.environ.get("LITELLM_TIMEOUT", "120")) # Match Go's 120s default
completion_prompt = os.environ.get(
"LITELLM_COMPLETION_PROMPT", "Say this is a test"
"LITELLM_COMPLETION_PROMPT", _DEFAULT_COMPLETION_PROMPT
)
embedding_text = os.environ.get(
"LITELLM_EMBEDDING_TEXT", "This is a test for vectorization."
"LITELLM_EMBEDDING_TEXT", _DEFAULT_EMBEDDING_TEXT
)
json_output = os.environ.get("LITELLM_JSON_OUTPUT", "").lower() == "true"
# Optional: only health-check these model IDs (comma-separated). E.g.:
@ -386,6 +425,7 @@ async def main():
timeout=timeout,
completion_prompt=completion_prompt,
embedding_text=embedding_text,
custom_auth_header=custom_auth_header,
)
# Load models from YAML if provided, otherwise fetch from API

74
scripts/health_check/health_check_client_README.md
View File

@ -30,6 +30,14 @@ export LITELLM_MODELS_YAML="/path/to/config.yaml"
python scripts/health_check/health_check_client.py
```
**Option 3: Use custom authentication header**
```bash
export LITELLM_BASE_URL="https://litellm.example.com"
export LITELLM_API_KEY="your-api-key"
export LITELLM_CUSTOM_AUTH_HEADER="x-custom-auth-header"
python scripts/health_check/health_check_client.py
```
### As a Docker Container
1. Build the Docker image:
@ -47,6 +55,16 @@ docker run --rm \
litellm/litellm-health-check:latest
```
3. Run with custom authentication header:
```bash
docker run --rm \
-e LITELLM_BASE_URL="https://litellm.example.com" \
-e LITELLM_API_KEY="your-api-key" \
-e LITELLM_CUSTOM_AUTH_HEADER="x-custom-auth-header" \
litellm/litellm-health-check:latest
```
### Parallel Execution (Stress Testing)
Run multiple health check containers in parallel:
@ -65,6 +83,30 @@ export LITELLM_API_KEY="your-api-key"
./scripts/health_check/run_parallel_health_checks.sh 16
```
**With Custom Auth Header:**
```powershell
$env:LITELLM_BASE_URL="https://litellm.example.com"
$env:LITELLM_API_KEY="your-api-key"
$env:LITELLM_CUSTOM_AUTH_HEADER="x-custom-auth-header"
.\scripts\health_check\run_parallel_health_checks.ps1 16
```
**With Custom Docker Image:**
```powershell
$env:LITELLM_BASE_URL="https://litellm.example.com"
$env:LITELLM_API_KEY="your-api-key"
$env:LITELLM_CUSTOM_AUTH_HEADER="x-custom-auth-header"
.\scripts\health_check\run_parallel_health_checks.ps1 -NumParallelJobs 16 -ImageName "your-registry/your-image:tag"
```
**Bash with Custom Image:**
```bash
export LITELLM_BASE_URL="https://litellm.example.com"
export LITELLM_API_KEY="your-api-key"
export LITELLM_CUSTOM_AUTH_HEADER="x-custom-auth-header"
./scripts/health_check/run_parallel_health_checks.sh 16 "your-registry/your-image:tag"
```
## Configuration
@ -73,14 +115,29 @@ export LITELLM_API_KEY="your-api-key"
- `LITELLM_BASE_URL` (required): Base URL of the LiteLLM proxy
- Example: `https://litellm.example.com`
- `LITELLM_API_KEY` (required): API key for authentication
- `LITELLM_CUSTOM_AUTH_HEADER` (optional): Custom header name for authentication
- Use this when your LiteLLM proxy uses a custom authentication header instead of the standard `Authorization` header
- Example: `x-custom-auth-header` (the API key will be sent as `Bearer <api_key>` in this header)
- `LITELLM_MODELS_YAML` (optional): Path to YAML config file with model_list
- If provided, reads models from YAML instead of fetching from API
- Example: `/path/to/config.yaml`
- `LITELLM_TIMEOUT` (optional): Request timeout in seconds (default: 120)
- `LITELLM_COMPLETION_PROMPT` (optional): Test prompt for chat/completion models (default: "Say this is a test")
- `LITELLM_EMBEDDING_TEXT` (optional): Test text for embedding models (default: "This is a test for vectorization.")
- `LITELLM_COMPLETION_PROMPT` (optional): Test prompt for chat/completion models (default: ~100k characters)
- `LITELLM_EMBEDDING_TEXT` (optional): Test text for embedding models (default: ~100k characters)
- `LITELLM_JSON_OUTPUT` (optional): Output results as JSON (default: false)
### Parallel Script Parameters
**PowerShell (`run_parallel_health_checks.ps1`):**
- `-NumParallelJobs` (optional): Number of parallel containers to run (default: 16)
- `-ImageName` (optional): Docker image to use (default: `litellm/litellm-health-check:latest`)
- `-ContainerRuntime` (optional): Container runtime to use (default: `docker`)
**Bash (`run_parallel_health_checks.sh`):**
- `[num_parallel_jobs]` (optional): Number of parallel containers to run (default: 16)
- `[image_name]` (optional): Docker image to use (default: `litellm/litellm-health-check:latest`)
- `[container_runtime]` (optional): Container runtime to use (default: `docker`)
## Output
### Standard Output (Human-Readable)
@ -166,7 +223,20 @@ Run multiple health checks in parallel:
**PowerShell:**
```powershell
# Using default image
.\scripts\health_check\run_parallel_health_checks.ps1 16
# Using custom image
.\scripts\health_check\run_parallel_health_checks.ps1 -NumParallelJobs 16 -ImageName "your-registry/your-image:tag"
```
**Bash:**
```bash
# Using default image
./scripts/health_check/run_parallel_health_checks.sh 16
# Using custom image
./scripts/health_check/run_parallel_health_checks.sh 16 "your-registry/your-image:tag"
```
### 3. CI/CD Integration

26
scripts/health_check/run_parallel_health_checks.ps1
View File

@ -49,6 +49,11 @@ Write-Host " 2. Set LITELLM_BASE_URL to the correct URL (e.g., http://host.do
Write-Host " 3. On Linux, you may need to use the host IP instead of host.docker.internal" -ForegroundColor Yellow
Write-Host ""
# Capture environment variables in parent scope for use in parallel block
$baseUrl = $env:LITELLM_BASE_URL
$apiKey = $env:LITELLM_API_KEY
$customAuthHeader = $env:LITELLM_CUSTOM_AUTH_HEADER
# Run parallel health checks
# This creates an infinite loop that keeps spawning containers
# Each container tests all models, then exits, and a new one starts
@ -57,13 +62,20 @@ while ($true) {
1..$NumParallelJobs | ForEach-Object -Parallel {
$runtime = $using:ContainerRuntime
$imageName = $using:ImageName
$baseUrl = $env:LITELLM_BASE_URL
$apiKey = $env:LITELLM_API_KEY
$baseUrl = $using:baseUrl
$apiKey = $using:apiKey
$customAuthHeader = $using:customAuthHeader
& $runtime run --rm `
-e LITELLM_BASE_URL="$baseUrl" `
-e LITELLM_API_KEY="$apiKey" `
-e LITELLM_JSON_OUTPUT="true" `
$imageName
$envVars = @(
"-e", "LITELLM_BASE_URL=$baseUrl",
"-e", "LITELLM_API_KEY=$apiKey",
"-e", "LITELLM_JSON_OUTPUT=true"
)
if ($customAuthHeader) {
$envVars += "-e", "LITELLM_CUSTOM_AUTH_HEADER=$customAuthHeader"
}
& $runtime run --rm $envVars $imageName
} -ThrottleLimit $NumParallelJobs
}

17
scripts/health_check/run_parallel_health_checks.sh
View File

@ -54,11 +54,18 @@ echo ""
# Function to run a single health check container
run_health_check() {
"$CONTAINER_RUNTIME" run --rm \
-e LITELLM_BASE_URL="$LITELLM_BASE_URL" \
-e LITELLM_API_KEY="$LITELLM_API_KEY" \
-e LITELLM_JSON_OUTPUT="true" \
"$IMAGE_NAME"
local env_vars=(
-e "LITELLM_BASE_URL=$LITELLM_BASE_URL"
-e "LITELLM_API_KEY=$LITELLM_API_KEY"
-e "LITELLM_JSON_OUTPUT=true"
)
# Pass through custom auth header if set
if [ -n "$LITELLM_CUSTOM_AUTH_HEADER" ]; then
env_vars+=(-e "LITELLM_CUSTOM_AUTH_HEADER=$LITELLM_CUSTOM_AUTH_HEADER")
fi
"$CONTAINER_RUNTIME" run --rm "${env_vars[@]}" "$IMAGE_NAME"
}
# Run parallel health checks