LLM Token and Dollar Budget Guardrail for Teams

id: llm-token-budget-cost-guardrail
namespace: company.team
description: |
  Enforce per-team LLM spend limits: every gateway usage event lands in Postgres,
  one query flags teams that reached their daily dollar cap, and over-budget teams
  get KV-quarantined, alerted in Slack, and filed as a GitHub issue.

inputs:
  - id: budgets
    type: JSON
    displayName: Per-team daily budgets
    description: One entry per team - team name, model filter (* matches all
      models), and the daily USD cap.
    defaults:
      - team: research
        model: "*"
        daily_limit_usd: 50
      - team: support
        model: "*"
        daily_limit_usd: 25
      - team: data
        model: gpt-4o
        daily_limit_usd: 15

  - id: repository
    type: STRING
    displayName: GitHub repository for budget issues
    description: owner/repo where over-budget teams get an issue, e.g. your-org/platform.
    defaults: your-org/your-repo

tasks:
  - id: record_usage
    type: io.kestra.plugin.jdbc.postgresql.Queries
    description: Persist the usage event posted by the LLM gateway or proxy.
    url: "{{ secret('POSTGRES_JDBC_URL') }}"
    username: "{{ secret('POSTGRES_USERNAME') }}"
    password: "{{ secret('POSTGRES_PASSWORD') }}"
    fetchType: NONE
    sql: |
      INSERT INTO llm_usage (team, model, prompt_tokens, completion_tokens, cost_usd, occurred_at)
      VALUES (
        '{{ trigger.body.team ?? "unknown" }}',
        '{{ trigger.body.model ?? "unknown" }}',
        {{ trigger.body.prompt_tokens ?? 0 }},
        {{ trigger.body.completion_tokens ?? 0 }},
        {{ trigger.body.cost_usd ?? 0 }},
        now()
      );

  - id: daily_spend_check
    type: io.kestra.plugin.jdbc.postgresql.Query
    description: Teams whose spend today reached their cap - budgets rendered as
      VALUES, joined against today's rollup.
    url: "{{ secret('POSTGRES_JDBC_URL') }}"
    username: "{{ secret('POSTGRES_USERNAME') }}"
    password: "{{ secret('POSTGRES_PASSWORD') }}"
    fetchType: FETCH
    sql: |
      WITH budgets(team, model, daily_limit_usd) AS (
        VALUES
        {% for b in inputs.budgets %}
          ('{{ b.team }}', '{{ b.model }}', {{ b.daily_limit_usd }}){% if not loop.last %},{% endif %}
        {% endfor %}
      ),
      spend AS (
        SELECT team, model,
               round(sum(cost_usd)::numeric, 2) AS spent_usd,
               sum(prompt_tokens + completion_tokens) AS tokens
        FROM llm_usage
        WHERE occurred_at >= date_trunc('day', now())
        GROUP BY team, model
      ),
      team_totals AS (
        SELECT team,
               round(sum(cost_usd)::numeric, 2) AS team_spent_usd,
               sum(prompt_tokens + completion_tokens) AS team_tokens
        FROM llm_usage
        WHERE occurred_at >= date_trunc('day', now())
        GROUP BY team
      )
      SELECT b.team,
             CASE WHEN b.model = '*' THEN 'all models' ELSE b.model END AS model,
             b.daily_limit_usd,
             CASE WHEN b.model = '*'
                  THEN coalesce(tt.team_spent_usd, 0)
                  ELSE coalesce(s.spent_usd, 0) END AS spent_usd,
             CASE WHEN b.model = '*'
                  THEN coalesce(tt.team_tokens, 0)
                  ELSE coalesce(s.tokens, 0) END AS tokens
      FROM budgets b
      LEFT JOIN team_totals tt ON b.model = '*' AND tt.team = b.team
      LEFT JOIN spend s ON b.model <> '*' AND s.team = b.team AND s.model = b.model
      WHERE CASE WHEN b.model = '*'
                 THEN coalesce(tt.team_spent_usd, 0)
                 ELSE coalesce(s.spent_usd, 0) END >= b.daily_limit_usd
      ORDER BY spent_usd DESC;

  - id: budget_gate
    type: io.kestra.plugin.core.flow.If
    description: Quarantine and notify per over-budget team, open a tracking issue,
      log an all-clear otherwise.
    condition: "{{ outputs.daily_spend_check.rows | length > 0 }}"
    then:
      - id: quarantine_offenders
        type: io.kestra.plugin.core.flow.Loop
        description: Set a KV budget-lock per offending team and alert its channel numbers.
        values: "{{ outputs.daily_spend_check.rows }}"
        tasks:
          - id: set_budget_lock
            type: io.kestra.plugin.core.kv.Set
            description: Store the breach row under llm_budget_lock/<team> so gateways or
              guard flows can refuse new calls.
            key: "llm_budget_lock/{{ item.value.team }}"
            value: "{{ item.value | json }}"
            kvType: JSON

          - id: alert_team
            type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
            url: "{{ secret('SLACK_WEBHOOK_URL') }}"
            messageText: |
              :money_with_wings: LLM daily budget hit for *{{ item.value.team }}* ({{ item.value.model }}): ${{ item.value.spent_usd }} of ${{ item.value.daily_limit_usd }} spent today ({{ item.value.tokens }} tokens). Route or pause new calls until tomorrow's reset.

      - id: open_issue
        type: io.kestra.plugin.github.issues.Create
        description: One labeled issue listing every offending team with the numbers.
        jwtToken: "{{ secret('GITHUB_TOKEN') }}"
        repository: "{{ inputs.repository }}"
        title: "LLM budget exceeded: {{ outputs.daily_spend_check.rows | length }}
          team(s) over daily cap"
        labels:
          - finops
          - llm-budget
        body: |
          ## Daily LLM budget breach

          | Team | Model | Spent (USD) | Cap (USD) | Tokens |
          |---|---|---|---|---|
          {% for r in outputs.daily_spend_check.rows %}
          | {{ r.team }} | {{ r.model }} | {{ r.spent_usd }} | {{ r.daily_limit_usd }} | {{ r.tokens }} |
          {% endfor %}

          Execution `{{ execution.id }}` recorded the breach at {{ execution.startDate }}.
          Each offending team is marked with a `llm_budget_lock/<team>` KV key until tomorrow's reset.

    else:
      - id: log_within_budget
        type: io.kestra.plugin.core.log.Log
        message: "All {{ inputs.budgets | length }} budgeted team(s) are within their
          daily LLM caps - {{ outputs.daily_spend_check.rows | length }}
          breach(es)."

outputs:
  - id: over_budget_count
    type: INT
    value: "{{ outputs.daily_spend_check.rows | length }}"

  - id: offenders
    type: JSON
    value: "{{ outputs.daily_spend_check.rows | json }}"

errors:
  - id: alert_on_failure
    type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
    url: "{{ secret('SLACK_WEBHOOK_URL') }}"
    messageText: ":rotating_light: LLM budget guard {{ flow.id }} FAILED (execution
      {{ execution.id }}). Spend is unmeasured while this fails. Error: {{
      error.message }}"

triggers:
  - id: usage_events
    type: io.kestra.plugin.core.trigger.Webhook
    description: Every completed LLM call posts its usage record here - wire your
      gateway or proxy to this URL. Replace the key with your own value.
    key: replace-with-a-random-key