Skip to content

Latest commit

 

History

History
118 lines (95 loc) · 4.64 KB

File metadata and controls

118 lines (95 loc) · 4.64 KB

Quotas & rate limiting (rag-quota)

Per-tenant quotas and rate limiting for AgentContextOS (Step 4.5). Enforcement flows through the existing PolicyEngine PDP (ADR-0005, ADR-0024).

Overview

Dimension Decision Shape Cap (TenantQuota)
QPS rate_limit sliding window (1 s) qps
Tokens / month quota_check sliding window (30 d) tokens_per_month
Cost / month quota_check sliding window (30 d), micro-dollars dollars_per_month
Queries / month quota_check sliding window (30 d) queries_per_month
Storage quota_check cumulative gauge storage_bytes

A denial raises RateLimitError (QPS) or QuotaExceededError (accruals), both HTTP 429 with a Retry-After header.

Usage

Build an enforcer

from rag_quota import (
    QuotaEnforcer, QuotaPolicy, QuotaRuntimeConfig, TenantQuotaLimits,
    InMemoryQuotaStore, QuotaPolicyEngine,
)
from rag_policy import NoopPolicyEngine

policy = QuotaPolicy(
    default_limits=TenantQuotaLimits(qps=50, tokens_per_month=5_000_000),
    overrides={"vip": TenantQuotaLimits(qps=500)},  # per-field merge with default
)
enforcer = QuotaEnforcer(
    store=InMemoryQuotaStore(),            # or RedisQuotaStore for cross-worker
    policy=policy,
    config=QuotaRuntimeConfig(micro_dollars_per_1k_tokens=500),  # $0.0005/1k
)
engine = QuotaPolicyEngine(enforcer, inner=NoopPolicyEngine())

rag.yaml

quotas:
  enabled: true            # opt-in — an enabled quota can reject requests
  fail_open: true          # store unreachable → admit + log quota.store_degraded
  qps_window_seconds: 1.0
  monthly_window_seconds: 2592000     # 30 days
  dollars_per_1k_tokens: 0.5          # drives the cost dimension (0 disables)
  default:                 # baseline; tenants[].quota overrides per-field
    qps: 50
    tokens_per_month: 5000000
    dollars_per_month: 250.0
    storage_bytes: 10737418240

tenants:
  - id: vip
    name: VIP
    quota:
      qps: 500             # qps overridden; other dims inherit `default`

When backends.cache.provider is redis/valkey, the gateway wires the RedisQuotaStore (cross-worker); otherwise it uses the in-process store.

CLI smoke test

ragctl quota --qps 3 --tokens 1000 --storage 4096 --price 0.5

Fires a burst at the QPS cap (last requests THROTTLE), advances a fake clock past the window, then exercises the monthly-token + storage caps and prints the usage snapshot. No infrastructure.

Status API

GET  /v1/status/quotas?tenant_id=acme        # one tenant's usage per dimension
POST /v1/status/quotas/{tenant}/reset        # clear usage (all dims, or ?dimension=qps)

GET resolves the tenant from tenant_id (query), else X-Tenant-Id (header), else "default". Both return empty / 404 when quotas are disabled.

Internals

  • QuotaStore SPI (rag_core.spi.quota_store) — check_and_consume, usage, adjust_gauge, gauge, reset, health. Keys are namespaced by ctx.tenant_id inside each implementation.
  • InMemoryQuotaStore — weighted two-slot sliding-window counter + gauge, injected clock (now) for deterministic tests.
  • RedisQuotaStore (rag_backends.quota.redis) — the same algorithm via one atomic Lua script per check (GET both window slots → estimate → conditional INCRBY), plus a gauge script. Raises on Redis failure (the enforcer owns the fail-open).
  • QuotaEnforcer — dimension dispatch, uncapped short-circuit, per-denial quota.exceeded event (PII-free), fail-open degrade, snapshot/reset.
  • QuotaPolicyEngine — answers rate_limit / quota_check; delegates read_chunk / ingest_doc / egress_text / execute_plan / filter_pushdown and health to the inner engine.

QuotaVerdict / QuotaSnapshot / QuotaDimension are frozen rag_core types; QuotaExceededError is a RateLimitError subtype.

Extension points

  • A different metering store — implement QuotaStore (e.g. a DynamoDB or Memcached backend) and inject it into QuotaEnforcer(store=…). Keep the weighted-window semantics so retry_after and conformance hold.
  • A real PDP — wrap your production PolicyEngine as the inner of QuotaPolicyEngine; quota decisions are answered here, everything else passes through to your engine.
  • Custom dimensionsQuotaSubject.kind is free-form; the enforcer maps the prefix (tokens / cost / queries / storage) to a QuotaDimension. Unknown kinds are allowed (never block on an unrecognised dimension).