{
  "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
}