{ "cells": [ { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "JRCXfhACct4Y" }, "source": [ "## User Based Rate Limiting Using LiteLLM\n", "- LiteLLM allows you to set budgets per user\n", "- Check if a given user has cross their allocated budget\n", "\n", "In this notebook we create a $0.0002 daily budget per user and make completion calls using the litellm budget manager" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "fl1kcLG8aaIV" }, "outputs": [], "source": [ "!pip install litellm uuid" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "zqRrpoQ3c6oQ" }, "source": [ "## Imports & Env variables" ] }, { "cell_type": "code", "execution_count": 13, "metadata": { "id": "CSkz8bmwdD3w" }, "outputs": [], "source": [ "import uuid\n", "import os\n", "os.environ['OPENAI_API_KEY'] = \"\"" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "ktqe3gSmdFQ4" }, "source": [ "## completion() with the budget manager\n", "\n", "This code does the following\n", "- Initializes a litellm.BudgetManager()\n", "- Checks if a budget exists for a user\n", " - Creates a $0.0002 budget if the user does not exisr\n", "- Makes a `litellm.completion()` request only if the user is under their budget" ] }, { "cell_type": "code", "execution_count": 23, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "pUN48YvmaiRU", "outputId": "082d6a8b-9aef-4794-9eac-7ba9823ea373" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "No budget exists for user: 29af95f8-c3c6-4c8c-b080-8b2d18d25432\n", "\n", "Creating a budget for user: 29af95f8-c3c6-4c8c-b080-8b2d18d25432, daily budget $0.0002\n", "\n", "User: 29af95f8-c3c6-4c8c-b080-8b2d18d25432 has spent $0, budget for user: $0.0002\n", "\n", "{\n", " \"id\": \"chatcmpl-7yAUkHQV8xdfldzzZnnnuVU8pl31b\",\n", " \"object\": \"chat.completion\",\n", " \"created\": 1694574378,\n", " \"model\": \"gpt-3.5-turbo-0613\",\n", " \"choices\": [\n", " {\n", " \"index\": 0,\n", " \"message\": {\n", " \"role\": \"assistant\",\n", " \"content\": \"Hello! I'm an AI, so I don't have emotions, but I'm here to assist you. How can I help you today?\"\n", " },\n", " \"finish_reason\": \"stop\"\n", " }\n", " ],\n", " \"usage\": {\n", " \"prompt_tokens\": 14,\n", " \"completion_tokens\": 29,\n", " \"total_tokens\": 43\n", " }\n", "}\n" ] }, { "data": { "text/plain": [ "{'status': 'success'}" ] }, "execution_count": 23, "metadata": {}, "output_type": "execute_result" } ], "source": [ "from litellm import BudgetManager, completion\n", "\n", "# Initializes a litellm.BudgetManager()\n", "budget_manager = BudgetManager(project_name=\"liteLLM_project\", client_type=\"hosted\") # see https://docs.litellm.ai/docs/budget_manager\n", "\n", "user_id = str(uuid.uuid4()) # create a new user id\n", "daily_budget = 0.0002\n", "\n", "# Checks if a budget exists for a user\n", "if not budget_manager.is_valid_user(user_id):\n", " # Creates a $0.0002 budget if the user does not exisr\n", " print(f\"No budget exists for user: {user_id}\\n\")\n", " print(f\"Creating a budget for user: {user_id}, daily budget ${daily_budget}\\n\")\n", " budget_manager.create_budget(total_budget=daily_budget, user=user_id, duration=\"daily\") # duration can be daily, weekly, monthly\n", "\n", "\n", "# Makes a `litellm.completion()` request only if the user is under their budget\n", "current_spend_for_user = budget_manager.get_current_cost(user=user_id)\n", "budget_for_user = budget_manager.get_total_budget(user_id)\n", "print(f\"User: {user_id} has spent ${current_spend_for_user}, budget for user: ${budget_for_user}\\n\")\n", "\n", "if current_spend_for_user <= budget_for_user:\n", " response = completion(model=\"gpt-3.5-turbo\", messages=[{\"role\": \"user\", \"content\": \"Hey, how's it going?\"}])\n", " budget_manager.update_cost(completion_obj=response, user=user_id)\n", "else:\n", " response = \"Sorry - no budget!\"\n", "\n", "print(response)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "yMOirNoBfmmc" }, "source": [ "## Make 10 calls to cross the budget per user\n", "- Code fails after user crossed their budget" ] }, { "cell_type": "code", "execution_count": 24, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "apKF3H-xbFXc", "outputId": "1c6ef0fe-e27e-4ead-adc6-2c7eb0214e44" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "User: 29af95f8-c3c6-4c8c-b080-8b2d18d25432 has spent $7.9e-05, budget for user: $0.0002\n", "\n", "User: 29af95f8-c3c6-4c8c-b080-8b2d18d25432 has spent $0.00015999999999999999, budget for user: $0.0002\n", "\n", "User: 29af95f8-c3c6-4c8c-b080-8b2d18d25432 has spent $0.00023899999999999998, budget for user: $0.0002\n", "\n", "User: 29af95f8-c3c6-4c8c-b080-8b2d18d25432 has exceeded budget, current spend $0.00023899999999999998, budget for user: $0.0002\n", "\n" ] } ], "source": [ "user_id = \"29af95f8-c3c6-4c8c-b080-8b2d18d25432\" # set in the previous cell\n", "\n", "for _ in range(10):\n", " # check if a given call can be made\n", " current_spend_for_user = budget_manager.get_current_cost(user=user_id)\n", " budget_for_user = budget_manager.get_total_budget(user_id)\n", " print(f\"User: {user_id} has spent ${current_spend_for_user}, budget for user: ${budget_for_user}\\n\")\n", " if current_spend_for_user <= budget_for_user:\n", " response = completion(model=\"gpt-3.5-turbo\", messages=[{\"role\": \"user\", \"content\": \"Hey, how's it going?\"}])\n", " budget_manager.update_cost(completion_obj=response, user=user_id)\n", " else:\n", " response = \"Sorry - no budget!\"\n", " print(f\"User: {user_id} has exceeded budget, current spend ${current_spend_for_user}, budget for user: ${budget_for_user}\\n\")\n", " break # no more requests\n", "\n", " # print(response)" ] } ], "metadata": { "colab": { "provenance": [] }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 0 }