Skip to main content

Building the Interception Script: Middleware for Local LLMs

Updated 7 min readMarkdown
Building the Interception Script: Middleware for Local LLMs

A guard model cannot defend a system if it is not placed directly in the path of network traffic. To utilize Llama Guard or Qwen3Guard effectively, you must construct an API gateway that intercepts all incoming traffic before it interacts with your primary generative model.

Using FastAPI in Python, you can implement a middleware script. Middleware acts as a literal "guard at the gate"—it catches the HTTP POST request, extracts the prompt, pauses the network request to run inference on the Guard model, and either routes it forward or throws a 403 Forbidden error.

Below is a complete, production-ready interception script. It assumes you have saved the Llama Guard logic from Article 1 into a file named guard.py.

from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import JSONResponse
import datetime
import time

# Import the guard functions defined in the previous scripts
from guard import moderate_input, moderate_output

app = FastAPI(title="AI Safety Gateway")

@app.middleware("http")
async def safety_interceptor(request: Request, call_next):
    """
    Middleware that intercepts all traffic to ensure safety before and after generation.
    """
    # Only intercept chat payload POST requests
    if request.method == "POST" and request.url.path == "/v1/chat":
        # 1. Read the incoming request body
        try:
            body_bytes = await request.body()
            body_json = await request.json()
            user_prompt = body_json.get("prompt", "")
        except Exception:
            return JSONResponse(status_code=400, content={"error": "Invalid JSON payload."})

        # 2. Log and Moderate the Input
        print(f"[{datetime.datetime.now()}] Intercepting prompt: {user_prompt[:50]}...")
        
        start_time = time.time()
        input_verdict = moderate_input(user_prompt)
        print(f"[Guard] Input evaluation took {time.time() - start_time:.2f}s. Verdict: {input_verdict}")

        if input_verdict.startswith("unsafe"):
            # Block the request completely; the main model is never triggered
            return JSONResponse(
                status_code=403, 
                content={"error": "Request blocked by safety policy.", "details": input_verdict}
            )

        # 3. Re-inject the body so the endpoint can read it (FastAPI consumes the stream)
        async def receive():
            return {"type": "http.request", "body": body_bytes}
        request._receive = receive

        # 4. Pass the safe request down to the main endpoint
        response = await call_next(request)

        # Note: In a true streaming application, output moderation requires intercepting 
        # the response chunks. For standard batch endpoints, you would capture the 
        # response body here, run moderate_output(), and rewrite the response if unsafe.
        
        return response

    # Pass through non-chat endpoints (like health checks) without moderation
    return await call_next(request)

@app.post("/v1/chat")
async def generate_text(payload: dict):
    """
    The main LLM endpoint. If traffic reaches here, the input was confirmed SAFE.
    """
    prompt = payload.get("prompt")
    
    # --- YOUR MAIN LLM INFERENCE CODE GOES HERE ---
    # Example: ai_response = main_model.generate(prompt)
    ai_response = f"This is a simulated AI response to: {prompt}"
    # ---------------------------------------------
    
    # Run output moderation before sending it back to the user
    output_verdict = moderate_output(prompt, ai_response)
    
    if output_verdict.startswith("unsafe"):
        return {"reply": "I am sorry, but I cannot fulfill this request as it violates safety guidelines."}
        
    return {"reply": ai_response}

How the FastAPI Interception Middleware Code Works

The middleware pattern acts as an HTTP proxy sitting directly in front of your LLM endpoints. Instead of modifying the main chat logic, it intercepts raw network requests before they touch the application and inspects network responses before they return to the client.

1. ASGI HTTP Request Interception (@app.middleware("http"))

FastAPI is built on top of Starlette and the ASGI (Asynchronous Server Gateway Interface) standard. The @app.middleware("http") decorator wraps around every incoming request. The request.url.path == "/v1/chat" filter ensures that non-generative endpoints (such as database health checks or metrics) bypass the guardrail logic entirely to prevent unnecessary compute overhead.

2. Stream Reading and Body Re-injection (request._receive)

In standard web frameworks, reading an incoming HTTP payload consumes the byte stream. Once Starlette reads await request.body(), the underlying stream buffer is empty. If passed down directly, the downstream endpoint (/v1/chat) would raise a 400 Bad Request or hang indefinitely trying to re-read an empty stream.

To solve this, the script re-injects the read bytes back into the request state using a synthetic receive function:

async def receive():
    return {"type": "http.request", "body": body_bytes}
request._receive = receive

This trick allows the guardrail to extract the prompt string while keeping the original request payload intact for FastAPI's Pydantic schema validation later down the pipeline.

3. Short-Circuiting Malicious Traffic

If moderate_input(user_prompt) returns unsafe, the middleware short-circuits the pipeline. It immediately returns a JSONResponse(status_code=403) directly to the client. The request is killed at the gateway stage; your primary text generation model (Qwen or Llama) is never invoked, saving GPU cycles and preventing adversarial context injection.

4. Post-Generation Output Filtering

If the prompt passes input moderation, execution proceeds via call_next(request), calling the main generation model. Once the main LLM completes its response, moderate_output() evaluates the (prompt, ai_response) pair. If the guardrail flags an unsafe response (such as a jailbreak bypass or unsafe hallucination), the raw output is replaced with a standard policy error message before reaching the user's browser.


Cons, Architectural Limitations, and Edge Cases

While implementing safety at the middleware level keeps code clean, it introduces severe operational challenges in production LLM environments.

1. The Streaming Response (SSE) Dilemma

The biggest flaw in basic middleware output moderation is that it breaks real-time token streaming.

  • Most modern AI interfaces stream text token-by-token to the client using Server-Sent Events (SSE) or WebSockets to minimize perceived latency.
  • For a traditional moderate_output() function to inspect a response, it must wait for the main LLM to finish generating the entire completion block.
  • If you buffer the entire text in middleware to moderate it, you strip away streaming, causing the user to sit staring at a loading spinner for several seconds.

Fixing this requires either token-by-token streaming classifiers (like Stream Qwen3Guard) or running output checks asynchronously while streaming chunks to the client with a kill-switch header.

2. Compound Latency Stacking (The 3x Penalty)

By sandwiching your main LLM between two full guardrail inference passes, you multiply latency:

$$\text{Total Latency} = \text{Guard}{\text{Input}} + \text{LLM}{\text{Generation}} + \text{Guard}_{\text{Output}}$$

If your guard model takes 300ms to read the input prompt and another 300ms to check the generated response, you add over half a second of static delay to every single request, regardless of whether the prompt was harmful or harmless.

3. VRAM and Compute Contention

Running FastAPI, your main generation model (e.g., Qwen-8B), and your guard model (e.g., Llama-Guard-3-8B) inside the same Python process or hardware instance creates massive GPU VRAM contention:

  • Loading two separate 8B model weights in bfloat16 requires approximately 32GB+ of VRAM just for base weights, before factoring in KV cache memory during concurrent requests.
  • If PyTorch blocks the main Python GIL (Global Interpreter Lock) during synchronous CUDA execution inside the middleware, it freezes FastAPI’s async event loop, causing request queuing and server timeouts under light load.

4. Payload Re-consumption Overhead and Memory Churn

Reading await request.body() into RAM and re-building synthetic _receive ASGI functions for every HTTP call introduces memory fragmentation in Python. On endpoints processing large context windows (e.g., 32k-token document uploads), copying multi-megabyte JSON strings multiple times in memory during middleware parsing leads to rapid memory spikes and elevated p99 latency under high request concurrency.


Architectural Summary Comparison

Operational Dimension In-Process Guardrail (Middleware) Microservice Architecture (Separate API)
Setup Complexity Simple single-file FastAPI script High (requires RPC / HTTP inter-service routing)
Streaming Support Breaks SSE unless buffered or patched Native (stream chunks through separate proxy worker)
VRAM Footprint Concentrated on one machine (VRAM bottleneck) Distributed across dedicated Guard & LLM GPUs
Fault Isolation High risk (Guard crash kills whole API) Isolated (Guard failure falls back to default rule)

About the author

No comments yet. Be the first to comment.