Flask AI API: LLM Inference Wrapper

You will build a Flask API that wraps a GPU-hosted LLM with proper endpoint design, streaming support, and Gunicorn deployment. By the end, you will have a production-ready inference wrapper on your dedicated GPU server that serves multiple clients through clean REST endpoints.

Project Setup

Install Flask and the supporting libraries. The API will proxy requests to a vLLM backend while adding validation, authentication, and error handling.

pip install flask gunicorn requests python-dotenv

# app.py
from flask import Flask, request, jsonify, Response
import requests as http_requests
import os

app = Flask(__name__)

VLLM_URL = os.getenv("VLLM_URL", "http://localhost:8000/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "meta-llama/Llama-3.1-8B-Instruct")
API_KEY = os.getenv("API_KEY", "your-secret-key")

Chat Completion Endpoint

Build a chat endpoint that validates input, forwards to vLLM, and returns a clean response.

@app.route("/api/v1/chat", methods=["POST"])
def chat():
    auth = request.headers.get("X-API-Key")
    if auth != API_KEY:
        return jsonify({"error": "Invalid API key"}), 403

    data = request.get_json()
    if not data or "messages" not in data:
        return jsonify({"error": "messages field required"}), 400

    payload = {
        "model": MODEL_NAME,
        "messages": data["messages"],
        "max_tokens": data.get("max_tokens", 256),
        "temperature": data.get("temperature", 0.7),
        "stream": False
    }

    try:
        resp = http_requests.post(f"{VLLM_URL}/chat/completions", json=payload, timeout=60)
        resp.raise_for_status()
        result = resp.json()
        return jsonify({
            "content": result["choices"][0]["message"]["content"],
            "model": result["model"],
            "tokens": result["usage"]["total_tokens"]
        })
    except http_requests.exceptions.Timeout:
        return jsonify({"error": "Inference timeout"}), 504
    except http_requests.exceptions.ConnectionError:
        return jsonify({"error": "Model server unavailable"}), 503

For the FastAPI alternative with async support, see the FastAPI inference server guide. For a comparison of both frameworks, check FastAPI vs Flask.

Streaming Endpoint

Stream tokens to the client as they are generated. Flask uses generator functions with the Response object.

@app.route("/api/v1/chat/stream", methods=["POST"])
def chat_stream():
    auth = request.headers.get("X-API-Key")
    if auth != API_KEY:
        return jsonify({"error": "Invalid API key"}), 403

    data = request.get_json()
    if not data or "messages" not in data:
        return jsonify({"error": "messages field required"}), 400

    payload = {
        "model": MODEL_NAME,
        "messages": data["messages"],
        "max_tokens": data.get("max_tokens", 512),
        "temperature": data.get("temperature", 0.7),
        "stream": True
    }

    def generate():
        with http_requests.post(f"{VLLM_URL}/chat/completions",
                                json=payload, stream=True, timeout=120) as resp:
            for line in resp.iter_lines():
                if line:
                    yield line.decode("utf-8") + "\n"
        yield "data: [DONE]\n\n"

    return Response(generate(), mimetype="text/event-stream")

Health Checks and Metrics

Add health and readiness endpoints for monitoring and load balancer integration.

import time
from collections import defaultdict

request_metrics = {"total": 0, "errors": 0, "latency_sum": 0}

@app.before_request
def track_start():
    request._start_time = time.time()

@app.after_request
def track_metrics(response):
    latency = time.time() - getattr(request, "_start_time", time.time())
    request_metrics["total"] += 1
    request_metrics["latency_sum"] += latency
    if response.status_code >= 400:
        request_metrics["errors"] += 1
    return response

@app.route("/health")
def health():
    try:
        resp = http_requests.get(f"{VLLM_URL}/models", timeout=5)
        return jsonify({"status": "healthy", "model": MODEL_NAME})
    except Exception:
        return jsonify({"status": "unhealthy"}), 503

@app.route("/metrics")
def metrics():
    avg_latency = (request_metrics["latency_sum"] / request_metrics["total"]
                   if request_metrics["total"] > 0 else 0)
    return jsonify({
        "total_requests": request_metrics["total"],
        "error_count": request_metrics["errors"],
        "avg_latency_ms": round(avg_latency * 1000, 2)
    })

For comprehensive monitoring, integrate with Prometheus and Grafana. For structured logging, see the ELK stack guide.

Gunicorn Deployment

Deploy Flask with Gunicorn for production. Configure workers based on your workload pattern.

# gunicorn.conf.py
bind = "0.0.0.0:8080"
workers = 4              # CPU workers for request handling
worker_class = "gthread"  # Threaded workers for I/O-bound proxy
threads = 2
timeout = 120            # Match inference timeout
keepalive = 5
accesslog = "-"
errorlog = "-"
loglevel = "info"

# Run
gunicorn -c gunicorn.conf.py app:app

# Or with Gevent for async-like behaviour
pip install gevent
gunicorn -w 4 -k gevent -b 0.0.0.0:8080 app:app

Use gthread workers when proxying to vLLM — each thread handles one HTTP connection to the backend while other threads accept new requests. For truly async patterns, the FastAPI approach is more natural.

Production Configuration

Place Flask behind Nginx for TLS and static asset handling. Use environment variables for all configuration.

# .env
VLLM_URL=http://localhost:8000/v1
MODEL_NAME=meta-llama/Llama-3.1-8B-Instruct
API_KEY=your-production-key
FLASK_ENV=production

For queue-based async processing instead of synchronous proxying, see the Redis queue guide. For an API gateway in front of your Flask server, check Kong/Traefik configuration. The self-hosting guide covers infrastructure planning, and our tutorials section has additional deployment patterns. Compare inference backends with vLLM vs Ollama.