Per-Team LLM Token Budgets at the Gateway

git clone https://github.com/tjorourke/solo-labs.git
cd solo/agentic-budgets-kind

export AGENTGATEWAY_LICENSE_KEY=...

./scripts/quick.sh up
./scripts/port-forward.sh   # leave running

async def chat_completions(request: Request) -> JSONResponse:
    body = await request.json()
    messages = body.get("messages") or []
    completion_text = random.choice(TEMPLATES)
    prompt_tokens     = _approx_prompt_tokens(messages)         # word_count × 1.3
    completion_tokens = random.randint(400, 1500)               # variable
    total_tokens      = prompt_tokens + completion_tokens

    return JSONResponse({
        "id": f"chatcmpl-{uuid.uuid4().hex[:24]}",
        "model": body.get("model") or "mock-essay-7b",
        "choices": [{"message": {"role": "assistant", "content": completion_text}}],
        "usage": {
            "prompt_tokens":     prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens":      total_tokens,    # ◄── the gateway debits this
        },
    })

claims := jwt.MapClaims{
    "iss":  "agentic-budgets-kind",
    "sub":  "dba",
    "team": "dba",
    "iat":  now.Unix(),
    "exp":  now.AddDate(10, 0, 0).Unix(),
    "aud":  "mock-llm",
}
tok := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
tok.Header["kid"] = "agentic-budgets-key-1"
signed, _ := tok.SignedString(priv)

apiVersion: enterpriseagentgateway.solo.io/v1alpha1
kind: EnterpriseAgentgatewayPolicy
metadata: { name: jwt-auth, namespace: agentgateway-system }
spec:
  targetRefs:
    - { group: gateway.networking.k8s.io, kind: Gateway, name: budgets-gateway }
  traffic:
    jwtAuthentication:
      mode: Strict
      providers:
        - issuer: "agentic-budgets-kind"
          audiences: ["mock-llm"]
          jwks:
            remote:
              jwksPath: "/.well-known/jwks.json"
              backendRef: { group: "", kind: Service, name: jwt-issuer, namespace: budgets, port: 8080 }
    transformation:
      request:
        set:
          - name: X-Team-ID
            value: jwt.team        # ← strict CEL expression, resolves to the team claim

apiVersion: ratelimit.solo.io/v1alpha1
kind: RateLimitConfig
metadata: { name: team-token-budget, namespace: agentgateway-system }
spec:
  raw:
    descriptors:
      - key: team
        value: "dba"
        rateLimit: { requestsPerUnit: 5000,  unit: HOUR }
      - key: team
        value: "support"
        rateLimit: { requestsPerUnit: 20000, unit: HOUR }
    rateLimits:
      - actions:
          - requestHeaders:
              descriptorKey: "team"
              headerName: "X-Team-ID"
        type: TOKEN

apiVersion: enterpriseagentgateway.solo.io/v1alpha1
kind: EnterpriseAgentgatewayPolicy
metadata: { name: team-budget, namespace: llm }   # same namespace as llm-route
spec:
  targetRefs:                                      # local refs — no namespace field
    - { group: gateway.networking.k8s.io, kind: HTTPRoute, name: llm-route }
  traffic:
    entRateLimit:
      global:
        rateLimitConfigRefs:
          - { name: team-token-budget, namespace: agentgateway-system }

async with httpx.AsyncClient(timeout=REQUEST_TIMEOUT) as client:
    resp = await client.post(LLM_URL, json=payload, headers={
        "Authorization": f"Bearer {LLM_JWT}",            # ◄── identity
        "Content-Type":  "application/json",
    })

if resp.status_code == 429:
    return (
        f"Sorry — your team's LLM token budget is exhausted. "
        f"Please try again later. (HTTP 429 from the gateway; "
        f"team={TEAM_LABEL!r}.)"
    )
…
return resp.json()["choices"][0]["message"]["content"]

apiVersion: enterpriseagentgateway.solo.io/v1alpha1
kind: EnterpriseAgentgatewayPolicy
metadata:
  name: token-attribution-logging
  namespace: agentgateway-system
spec:
  targetRefs:
    - group: gateway.networking.k8s.io
      kind: Gateway
      name: budgets-gateway
  frontend:
    accessLog:
      attributes:
        add:
          - { name: user_id,      expression: 'default(jwt.sub, "unknown")' }
          - { name: tokens_total, expression: 'default(string(llm.totalTokens), "n/a")' }
          - { name: tokens_in,    expression: 'default(string(llm.inputTokens), "n/a")' }
          - { name: tokens_out,   expression: 'default(string(llm.outputTokens), "n/a")' }
          - { name: model,        expression: 'default(llm.responseModel, "n/a")' }
          - { name: status,       expression: 'string(response.code)' }

kubectl --context kind-budgets -n agentgateway-system logs deploy/budgets-gateway \
  | grep -oE 'user_id=[^ ]+ .*tokens_total=[0-9]+' \
  | awk '{for(i=1;i<=NF;i++){if($i~/^user_id=/)u=substr($i,9);
          if($i~/^tokens_total=/)t=substr($i,14)} sum[u]+=t}
         END{for(u in sum) printf "%-16s %d tokens\n", u, sum[u]}'
# dba      18342 tokens
# support   9551 tokens

  frontend:
    accessLog:
      attributes:
        add: [ ... as above ... ]
      otlp:
        backendRef:
          kind: Service
          name: otel-collector
          namespace: monitoring
          port: 4317

# Then, per-user tokens in Loki:
#   sum by (user_id) (
#     sum_over_time({app="budgets-gateway"} | logfmt | unwrap tokens_total [1h])
#   )

# forward the gateway so curl can reach it
kubectl --context kind-budgets -n agentgateway-system \
  port-forward svc/budgets-gateway 8090:80 &

# fire 5 requests as each user, using their minted JWT
for u in alice bob carol support; do
  TOK=$(kubectl --context kind-budgets -n kagent \
    get secret jwt-$u -o jsonpath='{.data.token}' | base64 -d)
  for i in $(seq 1 5); do
    curl -s -o /dev/null -X POST http://localhost:8090/v1/chat/completions \
      -H "Authorization: Bearer $TOK" -H "Content-Type: application/json" \
      -d '{"model":"mock-essay-7b","messages":[{"role":"user","content":"write an essay about databases"}]}'
  done
  echo "sent 5 requests as $u"
done

# forward Grafana (port-forward.sh also does this)
kubectl --context kind-budgets -n monitoring \
  port-forward svc/monitoring-grafana 3000:80 &

# open http://localhost:3000  →  log in  admin / admin
# →  dashboard "Per-Team LLM Token Budgets"
# →  scroll to the "Tokens by user", "Per-user LLM access log",
#     and "Tokens per user over time" panels at the bottom