feat: production-level reliability, observability, and OpenAI compatibility improvements

.github/workflows/ci.yml

@@ -26,10 +26,19 @@ jobs:
       run: uv sync --all-extras --all-groups
 
     - name: Run Lint
-      run: uv run ruff check src
+      run: uv run ruff check .
 
     - name: Run Format Check
       run: uv run ruff format --check src
 
     - name: Run Tests
       run: uv run pytest -v -m "not integration"
+
+    - name: Dependency Audit
+      run: uv audit
+
+    - name: Gitleaks Secret Detection
+      uses: gitleaks/gitleaks-action@v2
+      env:
+        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+

CHANGELOG.md

@@ -8,6 +8,82 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Added
+- **Inference Concurrency Control with Semaphore Protection (`MAX_CONCURRENT_INFERENCES`)**:
+  - Implemented `asyncio.Semaphore` limit around neural network embedding and reranking inference to prevent GPU/CPU saturation and CUDA OOM crashes.
+  - Added configurable queue timeout (`INFERENCE_SEMAPHORE_TIMEOUT_SECONDS`, default 30s) returning `503 Service Unavailable` on sustained overload.
+  - Added concurrency safety test in `src/tests/test_concurrency_edge.py`.
+- **Client-Specific API Keys & Individual Rate Limits (`API_KEYS_MAP`)**:
+  - Supported multiple API keys via `API_KEYS` environment variable (comma-separated or JSON dictionary `{key: limit_per_minute}`).
+  - Applied constant-time `secrets.compare_digest` across all configured keys to prevent timing attacks.
+  - Integrated per-key custom rate limit overrides into `RateLimiter` middleware.
+  - Added unit test suites in `src/tests/test_auth.py` and `src/tests/test_rate_limit.py`.
+- **Production Kubernetes Deployment Manifests & Documentation (`deploy/kubernetes/`)**:
+  - Provided production-ready Kubernetes manifests: `deployment.yaml` (with liveness/readiness probes and Prometheus annotations), `service.yaml` (ClusterIP), and `hpa.yaml` (HorizontalPodAutoscaler scaling 1-5 pods based on CPU/Memory targets).
+  - Added deployment guide in `docs/deployment.md`.
+- **OpenAPI & Swagger Documentation Enhancements**:
+  - Added OpenAPI tags (`Embeddings`, `Reranking`, `Models`, `Health`, `Metrics`), endpoint summaries, descriptions, and standard response codes (400, 401, 413, 429, 503).
+  - Added realistic schema examples (`json_schema_extra`) for `EmbeddingRequest` and `RerankRequest`.
+  - Added unit test in `src/tests/test_extended_features.py`.
+- **Prometheus Metrics Instrumentation for Token Usage and Batch Size Distribution**:
+  - Added `http_prompt_tokens_total` Counter labeled by model to track total token consumption.
+  - Added `http_request_batch_size` Histogram labeled by endpoint with standard exponential buckets up to 256 items.
+  - Added unit test suite in `src/tests/test_metrics.py`.
+- **Model Warmup & Preloading on Application Startup (`PRELOAD_MODELS`)**:
+  - Implemented `PRELOAD_MODELS` environment variable supporting a comma-separated list of model names to load during application lifespan initialization.
+  - Eliminates first-request cold-start latency for production environments.
+  - Added unit test suite in `src/tests/test_config.py` and `src/tests/test_graceful_shutdown.py`.
+- **Graceful Shutdown & In-Flight Request Draining Middleware**:
+  - Implemented request task tracking and graceful drain during FastAPI application lifespan shutdown (`SHUTDOWN_DRAIN_TIMEOUT_SECONDS`, default: 10s).
+  - Automatically returns `503 Service Unavailable` with retry message for new incoming requests while shutting down.
+  - Added unit test suite in `src/tests/test_graceful_shutdown.py`.
+- **Configurable Precision & Mixed-Precision Inference (`TORCH_DTYPE`)**:
+  - Added `TORCH_DTYPE` configuration supporting `float16`, `bfloat16`, and `float32`.
+  - Integrated `torch.autocast` in multimodal model inference and passed `torch_dtype` to SentenceTransformer and CrossEncoder.
+  - Added unit test suite in `src/tests/test_torch_dtype.py`.
+- **OpenAI-Compatible `dimensions` & `encoding_format: "base64"` Support**:
+  - Added Matryoshka dimension truncation with automatic L2 re-normalization.
+  - Added IEEE 754 float32 little-endian Base64 embedding serialization (`encoding_format="base64"`).
+  - Applied formatting consistently across local PyTorch inference, TEI proxy, and multimodal embedding flows.
+  - Added unit test suite in `src/tests/test_dimensions_and_encoding.py`.
+- **IP / Token-based Rate Limiter Middleware (`429 Too Many Requests`)**:
+  - Implemented sliding-window token bucket rate limiter tracking requests per minute per IP/Bearer token (`RATE_LIMIT_PER_MINUTE`, default: 120).
+  - Included `Retry-After` header in 429 responses and automatically exempted internal health check/metric probes (`/health`, `/healthz`, `/ready`, `/metrics`).
+  - Added unit test suite in `src/tests/test_rate_limit.py`.
+- **OpenAI-Compatible Models Endpoint (`GET /v1/models`)**:
+  - Implemented standard OpenAI model listing endpoint returning all configured embedding and reranking models (`ModelList`, `ModelCard`).
+  - Added dedicated test suite `src/tests/test_models_endpoint.py`.
+- **HTTP Payload Size Limit Middleware (DoS / OOM Defense)**:
+  - Enforced 32MB maximum request body size (`MAX_PAYLOAD_SIZE`), returning `413 Payload Too Large` for oversized requests before parsing.
+  - Added dedicated test suite `src/tests/test_payload_limit.py`.
+- **Dynamic Model Unloading & Memory Reclamation (`POST /v1/models/unload`)**:
+  - Added endpoint to dynamically unload specific or all models from cache, triggering `torch.cuda.empty_cache()` and garbage collection to free RAM/VRAM.
+  - Added dedicated test suite `src/tests/test_model_unload.py`.
+- **Structured JSON Logging with Request ID Tracking**:
+  - Integrated JSON log formatting and `X-Request-ID` correlation via ContextVars in HTTP middleware.
+  - Automatically captures HTTP method, endpoint, status code, and latency in standard JSON output for APM/log aggregation.
+  - Added dedicated test suite in `src/tests/test_logger.py`.
+- **Prometheus Metrics Instrumentation (`/metrics`)**:
+  - Integrated `prometheus_client` exposing standard Prometheus metrics for HTTP request count and latency histograms with endpoint grouping.
+  - Added dedicated test suite `src/tests/test_metrics.py`.
+- **Readiness Probe Endpoint (`/ready`)**:
+  - Added dedicated `/ready` endpoint verifying GPU availability and loaded model cache keys, decoupling readiness from liveness (`/health`, `/healthz`).
+  - Added unit test suite in `src/tests/test_ready.py`.
+- **CI / CD Automated Auditing & Secret Scanning**:
+  - Enforced `uv audit` and `gitleaks` in GitHub Actions CI workflow (`.github/workflows/ci.yml`) per `.rules/ci.md`.
+  - Expanded `ruff check` to entire repository (`.`).
+
+### Changed
+- **Multimodal Batch Inference Optimization**:
+  - Refactored `VisualizedBGEEmbeddingModel.encode_multimodal` in `src/app/models.py` to batch preprocessed image tensors and tokenized text together instead of processing items sequentially, drastically improving multi-item inference throughput.
+- **Asynchronous TEI Proxy**:
+  - Upgraded TEI proxy client in `src/app/main.py` to use `httpx.AsyncClient` with pooled connections, and converted `_proxy_to_tei` and service callers to full `async/await` execution to eliminate event loop blocking.
+
+### Fixed
+- **DNS Rebinding & TOCTOU Mitigation in Multimodal Image Downloads**:
+  - Implemented `SafeNetworkBackend` in `src/app/image_utils.py` with custom `httpcore.AsyncNetworkBackend` that pins the TCP connection target to the validated, safe IP address resolved during SSRF validation while retaining the original Host and SNI headers.
+  - Eliminated the vulnerability window between DNS resolution and HTTP stream connection.
+
+
 - **Multimodal (Diagram + Text) Full Support**:
   - Integrated `bge-visualized-m3` model for composite image + text and image-only embeddings in 1024 dimensions.
   - Added support for Flat schema (`FlatMultimodalItem`) and OpenAI ContentPart format (`[{"type": "text"}, {"type": "image_url"}]`).

Dockerfile

@@ -63,7 +63,8 @@ EXPOSE 8000
 ENV GUNICORN_WORKERS=2
 
 HEALTHCHECK --interval=10s --timeout=5s --start-period=20s --retries=3 \
-  CMD python3 -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8000/health', timeout=3)" || exit 1
+  CMD python3 -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8000/healthz', timeout=3)" || exit 1
+
 
 # Command to run the application using Gunicorn from the virtual environment
 CMD ["sh", "-c", "gunicorn --workers ${GUNICORN_WORKERS} --worker-class uvicorn.workers.UvicornWorker --bind 0.0.0.0:8000 --timeout 600 --worker-tmp-dir /dev/shm --keep-alive 5 src.app.main:app"]

Dockerfile.cpu

@@ -54,7 +54,8 @@ EXPOSE 8000
 ENV GUNICORN_WORKERS=2
 
 HEALTHCHECK --interval=10s --timeout=5s --start-period=20s --retries=3 \
-  CMD python3 -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8000/health', timeout=3)" || exit 1
+  CMD python3 -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8000/healthz', timeout=3)" || exit 1
+
 
 # Command to run the application using Gunicorn from the virtual environment
 CMD ["sh", "-c", "gunicorn --workers ${GUNICORN_WORKERS} --worker-class uvicorn.workers.UvicornWorker --bind 0.0.0.0:8000 --timeout 120 --worker-tmp-dir /dev/shm --keep-alive 5 src.app.main:app"]

deploy/kubernetes/deployment.yaml

@@ -0,0 +1,48 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: openai-compatible-api
+  labels:
+    app: openai-compatible-api
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      app: openai-compatible-api
+  template:
+    metadata:
+      labels:
+        app: openai-compatible-api
+      annotations:
+        prometheus.io/scrape: "true"
+        prometheus.io/path: "/metrics"
+        prometheus.io/port: "8000"
+    spec:
+      containers:
+      - name: api
+        image: openai-compatible-api:latest
+        imagePullPolicy: IfNotPresent
+        ports:
+        - containerPort: 8000
+          name: http
+        resources:
+          requests:
+            cpu: "1"
+            memory: "2Gi"
+          limits:
+            cpu: "2"
+            memory: "4Gi"
+        livenessProbe:
+          httpGet:
+            path: /healthz
+            port: http
+          initialDelaySeconds: 20
+          periodSeconds: 10
+          timeoutSeconds: 5
+        readinessProbe:
+          httpGet:
+            path: /ready
+            port: http
+          initialDelaySeconds: 10
+          periodSeconds: 10
+          timeoutSeconds: 5

deploy/kubernetes/hpa.yaml

@@ -0,0 +1,26 @@
+apiVersion: autoscaling/v2
+kind: HorizontalPodAutoscaler
+metadata:
+  name: openai-compatible-api
+  labels:
+    app: openai-compatible-api
+spec:
+  scaleTargetRef:
+    apiVersion: apps/v1
+    kind: Deployment
+    name: openai-compatible-api
+  minReplicas: 1
+  maxReplicas: 5
+  metrics:
+    - type: Resource
+      resource:
+        name: cpu
+        target:
+          type: Utilization
+          averageUtilization: 75
+    - type: Resource
+      resource:
+        name: memory
+        target:
+          type: Utilization
+          averageUtilization: 80

deploy/kubernetes/service.yaml

@@ -0,0 +1,15 @@
+apiVersion: v1
+kind: Service
+metadata:
+  name: openai-compatible-api
+  labels:
+    app: openai-compatible-api
+spec:
+  selector:
+    app: openai-compatible-api
+  ports:
+    - protocol: TCP
+      port: 8000
+      targetPort: 8000
+      name: http
+  type: ClusterIP

docker-compose.yml

@@ -20,7 +20,8 @@ services:
     volumes:
       - ./.cache/models:/home/appuser/.cache/huggingface
     healthcheck:
-      test: ["CMD", "python3", "-c", "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8000/health', timeout=3)"]
+      test: ["CMD", "python3", "-c", "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8000/healthz', timeout=3)"]
+
       interval: 10s
       timeout: 5s
       retries: 3
@@ -55,7 +56,8 @@ services:
               count: all
               capabilities: [ gpu ]
     healthcheck:
-      test: ["CMD", "python3", "-c", "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8000/health', timeout=3)"]
+      test: ["CMD", "python3", "-c", "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8000/healthz', timeout=3)"]
+
       interval: 10s
       timeout: 5s
       retries: 3

docs/deployment.md

@@ -0,0 +1,54 @@
+# Kubernetes Deployment
+
+This document provides instructions on how to deploy the OpenAI-Compatible API using Kubernetes.
+
+The production-ready manifests are located in the `deploy/kubernetes` directory and include a Deployment, a Service, and a Horizontal Pod Autoscaler (HPA).
+
+## Prerequisites
+
+- A running Kubernetes cluster.
+- `kubectl` configured to interact with your cluster.
+- A metrics server installed in the cluster (required for HPA to work).
+
+## Deployment Manifests
+
+The `deploy/kubernetes` directory contains three manifests:
+
+1. **`deployment.yaml`**: Contains the main `Deployment` object.
+    - Sets up the `openai-compatible-api` container running on port 8000.
+    - Defines a `readinessProbe` checking the `/ready` endpoint to ensure the application only receives traffic when fully loaded.
+    - Defines a `livenessProbe` checking the `/healthz` endpoint to restart pods if they become unresponsive.
+    - Specifies CPU and Memory `requests` and `limits` to ensure efficient scheduling and resource management.
+    - Includes annotations (`prometheus.io/scrape: "true"`, `prometheus.io/path: "/metrics"`, `prometheus.io/port: "8000"`) to allow Prometheus to automatically scrape metrics.
+2. **`service.yaml`**: Exposes the `Deployment` on port 8000 via a `ClusterIP` Service.
+3. **`hpa.yaml`**: Contains the `HorizontalPodAutoscaler` configuration.
+    - Automatically scales the number of pods between 1 and 5 based on target CPU (75%) and memory (80%) utilization.
+
+## Deploying to Kubernetes
+
+To deploy the application to your cluster, apply the manifests using `kubectl`:
+
+```bash
+kubectl apply -f deploy/kubernetes/deployment.yaml
+kubectl apply -f deploy/kubernetes/service.yaml
+kubectl apply -f deploy/kubernetes/hpa.yaml
+```
+
+Alternatively, you can apply the entire directory at once:
+
+```bash
+kubectl apply -f deploy/kubernetes/
+```
+
+## Monitoring
+
+- **Status**: Check the status of your pods to ensure they are running successfully:
+  ```bash
+  kubectl get pods -l app=openai-compatible-api
+  ```
+- **Autoscaling**: Verify the Horizontal Pod Autoscaler is correctly fetching metrics:
+  ```bash
+  kubectl get hpa openai-compatible-api
+  ```
+  *(Note: It may take a few minutes for the HPA to collect metrics after initial deployment.)*
+- **Metrics**: If Prometheus is installed in your cluster and configured to honor scrape annotations, it will automatically begin scraping the `/metrics` endpoint on port 8000 of the deployed pods.