[Feature]: Allow key budgets to be reset

[krrishdholakia]

The Feature

Budgets are set per key, will reset over a time period

Motivation, pitch

allow app owners (i.e. key) to have budgets that reset after a time period

Twitter / LinkedIn details

No response

[krrishdholakia]

Tentative backend code:

Install APScheduler:

pip install apscheduler

Edit your main.py to add budget reset logic.

from apscheduler.schedulers.asyncio import AsyncIOScheduler

scheduler = AsyncIOScheduler()

@scheduler.scheduled_job('interval', days=30)
async def reset_budgets():
    now = datetime.datetime.utcnow()
    keys_to_reset = await prisma.apiKey.find_many(
        where={
            'resetAt': {
                'lt': now
            },
        }
    )

    for key in keys_to_reset:
        key.budget = initial_budget
        key.resetAt = now + datetime.timedelta(days=30)
        await prisma.apiKey.update({
            'where': {'id': key.id},
            'data': key
        })

# inside your startup event after prisma.connect()
scheduler.start()

[krrishdholakia]

We'll need to introduce 1 new field into /key/generate:

• budget_duration: Specify the length of time the budget is valid for. If null, default is set to 1 hour. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d").

2 new fields in prisma.schema:

• budget_duration
• budget_reset_at # <- internal value, used to identify if budget needs to be reset

[krrishdholakia]

Thinking

Sample curl

curl 'http://0.0.0.0:8000/key/generate' \
--header 'Authorization: Bearer <your-master-key>' \
--header 'Content-Type: application/json' \
--data-raw '{
  "team_id": "core-infra", # [OPTIONAL]
  "max_budget": 0.0005,
  "budget_duration": "10s" 
}'

[krrishdholakia]

PR created - #1570

[krrishdholakia]

PR Merged and passing testing.