llm-guard-gateway

A governance gateway for LLM and agent traffic that blocks prompt-injection, redacts PII, enforces per-tenant token budgets, and serves repeat questions from a semantic cache, cutting served p50 latency from 44.2 ms to 4.6 ms in the included benchmark.

What this solves

• Teams wiring LLMs into products re-pay for near-duplicate prompts; the semantic cache collapses them onto one stored completion, and the benchmark measured an 87.7 percent hit rate at a 70 percent duplicate ratio.
• Prompt-injection reaches the model before any control does in most stacks; here a noisy-OR firewall scores every prompt in microseconds and blocked the classic override-plus-exfiltration attack at score 0.82 before a single token was spent.
• One tenant can silently burn a shared OpenAI budget; a continuously-refilled per-tenant token bucket sheds over-quota traffic with an explicit blocked_budget decision instead of a surprise invoice.

Why this exists

An enterprise platform team fronting Azure OpenAI for many internal apps has three recurring costs: duplicate spend (the same questions asked thousands of times in different words), security exposure (prompts that try to hijack the system prompt of downstream copilots), and unbounded consumption (no per-team metering). At list pricing for GPT-4-class models, a workload of one million requests a month with the 70 percent near-duplicate ratio used in this benchmark pays for roughly 700,000 completions it already produced. These are representative figures, not customer data; the duplicate ratio is a benchmark parameter you can change in one flag.

The gateway is a single enforcement point in front of the model provider. Every request passes a fixed, auditable pipeline: an injection firewall (compiled-regex signals combined noisy-OR, so independent weak signals compound), a per-tenant token-budget reservation, typed PII redaction, then a semantic cache lookup that embeds the prompt and serves any stored completion within 0.92 cosine similarity. Only a miss reaches the model. Every decision emits a structured audit event containing the redacted prompt metadata, never raw PII. The same capability is exposed twice: as a REST endpoint and as a Model Context Protocol tool (guard_prompt), so AI agents get the guardrails by speaking MCP instead of calling the provider directly.

Measured on the included load test (3,000 requests per level, 70 percent duplicate prompts, 2 vCPU container, 40 ms simulated model latency): enabling the cache moved p50 from 44.23 ms to 4.62 ms and throughput from 223 to 620 requests per second at 10 concurrent clients, with an 70.6 to 87.7 percent measured hit rate across levels. Raw results are in benchmark/results/.

Architecture

The local profile (default) swaps every backing service for an in-memory adapter behind the same interface, so the entire system, tests, and benchmark run offline with zero credentials. The prod profile binds the same interfaces to Redis (distributed budget), Postgres + pgvector (cache store), Kafka (audit events), and Azure OpenAI (embeddings and completions).

Live demo

Real requests against the running gateway, showing a cache miss, a sub-millisecond semantic hit with the identical completion, a blocked injection attempt, and the Prometheus counters that result:

API

Tech stack

Technology	Role in this project	Why chosen here
Python 3.11 + FastAPI	Gateway service and OpenAPI surface	Async-first: the request path is IO-shaped (model call) and benefits from cooperative concurrency under load
NumPy	Vector store similarity scan	One BLAS matrix-vector product gives exact cosine over the 10k-entry working set in microseconds; no ANN dependency
pydantic + pydantic-settings	Request validation and env-driven config	Rejects malformed input at the boundary (422 before any guard runs); profile switch is one env var
structlog	Structured JSON logging	Every decision is one JSON object with bound request context; drops into Splunk or Azure Monitor without parsing
pgvector (prod adapter)	Cache store beyond one process	Keeps the cache in Postgres, which the platform already operates, instead of adding a vector database
Redis (prod adapter)	Distributed token budget	Atomic Lua-script bucket shared across gateway replicas
Kafka (prod adapter)	Audit event stream	Fire-and-forget producer keeps compliance persistence off the request path
MCP (JSON-RPC 2.0)	Agent-facing tool surface	Agents route model calls through the same guardrails instead of around them
Docker + compose	Production-shaped local stack	One command brings up gateway, pgvector, Redis, and Kafka wired together
pytest + pytest-cov + ruff, GitHub Actions	Tests, coverage gate at 90, lint	CI fails on lint or a coverage drop; measured coverage is 98 percent

Quickstart

Prerequisites: Python 3.11+, git. Docker only for the prod-shaped stack.

git clone https://github.com/<you>/llm-guard-gateway.git
cd llm-guard-gateway
python -m venv .venv && source .venv/bin/activate
pip install -e ".[dev]"

# run the test suite
pytest --cov=llm_guard_gateway

# start the gateway (local profile: fully offline, in-memory adapters)
uvicorn llm_guard_gateway.main:app --port 8080

# exercise it
curl -s localhost:8080/v1/guard -H 'content-type: application/json' \
  -d '{"tenant":"team-a","prompt":"summarize the incident report"}'
curl -s localhost:8080/metrics

Production-shaped stack (gateway + pgvector + Redis + Kafka):

cp .env.example .env   # fill in Azure OpenAI values for real completions
docker compose up -d

Load test against a running gateway:

python benchmark/loadtest.py --url http://localhost:8080 \
  --requests 3000 --concurrency 10 20 40 --duplicate-ratio 0.7 --label myrun

Performance under load

Methodology: closed-loop async load generator (benchmark/loadtest.py, httpx), 3,000 requests per concurrency level with a fresh prompt set per level, 70 percent duplicate ratio against a 6-template pool. Environment: 2 vCPU Linux container, local profile, 40 ms simulated model latency, token budget raised so the rate limiter does not shed the single benchmark tenant (isolating cache and pipeline behaviour). Raw JSON in benchmark/results/.

Cache	Concurrency	Throughput (rps)	Hit rate	p50 (ms)	p95 (ms)	p99 (ms)
off	10	223.2	0.0	44.23	48.87	53.14
off	20	441.7	0.0	44.14	50.67	71.19
off	40	153.2	0.0	77.08	1180.83	2187.37
on	10	620.0	70.6%	4.62	46.09	50.51
on	20	416.1	80.9%	29.82	140.30	204.68
on	40	255.0	87.7%	95.40	486.71	790.84

Where it degrades: on 2 vCPUs the event loop saturates at 40 concurrent clients; the uncached p99 blows out to 2.19 s and one request errored, because every request holds a 40 ms model slot while new work keeps arriving. The cache softens the knee (p99 0.79 s at the same load) but does not remove it; the real remedy is horizontal replicas behind a load balancer, which the stateless design and the Redis budget adapter exist to allow. Tail latencies at concurrency 10 with cache on still show ~46 ms at p95 because 29 percent of requests are misses that pay the full model latency.

Architecture decisions

Two ADRs in docs/adr/: ADR-001 (exact cosine scan over ANN index, pgvector as the scale-out path) and ADR-002 (the boring choice: regex noisy-OR firewall over an LLM judge, and why the judge is itself an injection target).

Intentionally out of scope

• Streaming responses. The cache stores complete completions; streaming needs chunk-level storage and replay. Add when a consumer actually requires server-sent events.
• Response-side content filtering. The gateway governs what goes to the model, not what comes back. Add an output guard stage if the model output is user-facing rather than developer-facing.
• LRU or frequency-weighted cache eviction. FIFO is deliberate simplicity; swap the eviction policy when hit-rate telemetry shows hot entries being churned out, not before.

Security and compliance

Secrets come only from environment variables locally (.env is gitignored; .env.example documents every key) and from Azure Key Vault via managed identity in the production path. Raw prompts containing PII are redacted before they are cached, sent to the model, or written to any log or audit event; the audit stream carries redaction counts, never the original values. The MCP surface exposes exactly one tool with a validated schema. CI runs lint and tests on every push; the container image is multi-stage with only runtime dependencies in the final layer.

Failure modes

Failure	Detection	Behaviour	Recovery
Model provider down or rate-limiting	HTTP errors surface in structured logs and llmguard_llm_calls_total stalls	Cache hits keep serving; misses fail fast with the provider error rather than queueing	Provider retry with backoff belongs in the model client adapter; cached traffic rides through the outage
Redis (prod budget) unavailable	Connection errors on try_consume	Fail closed for budget enforcement is configurable; local in-process bucket is the degraded fallback	Reconnect; buckets refill from wall clock, no state to rebuild
Postgres/pgvector (prod cache) down	Lookup errors	Treat every request as a miss: correctness preserved, cost and latency rise	Cache repopulates organically on recovery; no warm-up job needed
Kafka audit broker down	Producer errors in logs	Fire-and-forget emit fails without blocking the request path; events are lost, decisions still logged locally	Restore broker; if audit is compliance-critical, switch the sink to an outbox table
Poisoned cache entry (bad completion stored)	Consumer reports; entry is traceable by key in the audit stream	Entry serves until evicted	FIFO bound caps exposure at 10k entries; a delete-by-key admin endpoint is listed in Future Work

Hardest problem solved

The semantic cache returned a completely wrong answer during integration testing: a prompt asking to reverse a linked list was served the cached summary of an earnings report. Cache hit rate looked excellent; correctness was silently broken. The kind of bug that ships.

Diagnosis started from a failing test I wrote to pin the cache's contract (test_unrelated_prompt_does_not_hit) and a direct unit test of the vector store's scoring. That second test made the root cause obvious: search scored candidates with a raw dot product, but the hashing embedder returns unnormalized term-frequency vectors, so magnitude scales with prompt length. A long prompt's dot product against anything could exceed the 0.92 threshold on magnitude alone, regardless of direction. The similarity threshold was meaningless.

The fix (baaff85) L2-normalizes vectors on insert and at query time, making the score true cosine similarity, bounded and scale-invariant, with regression tests asserting the self/orthogonal/scaled cases and that unrelated prompts miss. ADR-001 records the design consequence: similarity thresholds are only interpretable if the metric is actually cosine.

Future work

• Async embedding-similarity injection detection as a second, non-blocking layer that mines new regex rules from near-miss traffic (ADR-002 lays out why it must not block).
• Reconcile tokens_charged with the provider's actual usage on response, refunding the difference to the tenant's bucket.
• Admin endpoints: delete-by-key cache invalidation and per-tenant budget inspection.
• OpenTelemetry trace propagation through the pipeline stages so a slow request shows which stage paid the latency.
• First metric to watch in production: cache hit rate per tenant. It is the whole cost case; if it sits under 20 percent for a tenant, their prompts carry volatile context (timestamps, ids) and need prompt-normalization before the gateway can help them.

License

MIT