{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "680oRk1af-xJ"
      },
      "source": [
        "# Environment Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X7TgJFn8f88p"
      },
      "outputs": [],
      "source": [
        "import csv\n",
        "from typing import Optional\n",
        "import httpx, json\n",
        "import asyncio\n",
        "\n",
        "proxy_base_url = \"http://0.0.0.0:4000\" # 👈 SET TO PROXY URL\n",
        "master_key = \"sk-1234\" # 👈 SET TO PROXY MASTER KEY"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rauw8EOhgBz5"
      },
      "outputs": [],
      "source": [
        "## GLOBAL HTTP CLIENT ## - faster http calls\n",
        "class HTTPHandler:\n",
        "    def __init__(self, concurrent_limit=1000):\n",
        "        # Create a client with a connection pool\n",
        "        self.client = httpx.AsyncClient(\n",
        "            limits=httpx.Limits(\n",
        "                max_connections=concurrent_limit,\n",
        "                max_keepalive_connections=concurrent_limit,\n",
        "            )\n",
        "        )\n",
        "\n",
        "    async def close(self):\n",
        "        # Close the client when you're done with it\n",
        "        await self.client.aclose()\n",
        "\n",
        "    async def get(\n",
        "        self, url: str, params: Optional[dict] = None, headers: Optional[dict] = None\n",
        "    ):\n",
        "        response = await self.client.get(url, params=params, headers=headers)\n",
        "        return response\n",
        "\n",
        "    async def post(\n",
        "        self,\n",
        "        url: str,\n",
        "        data: Optional[dict] = None,\n",
        "        params: Optional[dict] = None,\n",
        "        headers: Optional[dict] = None,\n",
        "    ):\n",
        "        try:\n",
        "            response = await self.client.post(\n",
        "                url, data=data, params=params, headers=headers\n",
        "            )\n",
        "            return response\n",
        "        except Exception as e:\n",
        "            raise e\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7LXN8zaLgOie"
      },
      "source": [
        "# Import Sheet\n",
        "\n",
        "\n",
        "Format: | ID | Name | Max Budget |"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oiED0usegPGf"
      },
      "outputs": [],
      "source": [
        "async def import_sheet():\n",
        "    tasks = []\n",
        "    http_client = HTTPHandler()\n",
        "    with open('my-batch-sheet.csv', 'r') as file:\n",
        "        csv_reader = csv.DictReader(file)\n",
        "        for row in csv_reader:\n",
        "            task = create_user(client=http_client, user_id=row['ID'], max_budget=row['Max Budget'], user_name=row['Name'])\n",
        "            tasks.append(task)\n",
        "            # print(f\"ID: {row['ID']}, Name: {row['Name']}, Max Budget: {row['Max Budget']}\")\n",
        "\n",
        "    keys = await asyncio.gather(*tasks)\n",
        "\n",
        "    with open('my-batch-sheet_new.csv', 'w', newline='') as new_file:\n",
        "        fieldnames = ['ID', 'Name', 'Max Budget', 'keys']\n",
        "        csv_writer = csv.DictWriter(new_file, fieldnames=fieldnames)\n",
        "        csv_writer.writeheader()\n",
        "\n",
        "        with open('my-batch-sheet.csv', 'r') as file:\n",
        "            csv_reader = csv.DictReader(file)\n",
        "            for i, row in enumerate(csv_reader):\n",
        "                row['keys'] = keys[i]  # Add the 'keys' value from the corresponding task result\n",
        "                csv_writer.writerow(row)\n",
        "\n",
        "    await http_client.close()\n",
        "\n",
        "asyncio.run(import_sheet())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E7M0Li_UgJeZ"
      },
      "source": [
        "# Create Users + Keys\n",
        "\n",
        "- Creates a user\n",
        "- Creates a key with max budget"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NZudRFujf7j-"
      },
      "outputs": [],
      "source": [
        "\n",
        "async def create_key_with_alias(client: HTTPHandler, user_id: str, max_budget: float):\n",
        "    global proxy_base_url\n",
        "    if not proxy_base_url.endswith(\"/\"):\n",
        "        proxy_base_url += \"/\"\n",
        "    url = proxy_base_url + \"key/generate\"\n",
        "\n",
        "    # call /key/generate\n",
        "    print(\"CALLING /KEY/GENERATE\")\n",
        "    response = await client.post(\n",
        "        url=url,\n",
        "        headers={\"Authorization\": f\"Bearer {master_key}\"},\n",
        "        data=json.dumps({\n",
        "            \"user_id\": user_id,\n",
        "            \"key_alias\": f\"{user_id}-key\",\n",
        "            \"max_budget\": max_budget # 👈 KEY CHANGE: SETS MAX BUDGET PER KEY\n",
        "        })\n",
        "    )\n",
        "    print(f\"response: {response.text}\")\n",
        "    return response.json()[\"key\"]\n",
        "\n",
        "async def create_user(client: HTTPHandler, user_id: str, max_budget: float, user_name: str):\n",
        "    \"\"\"\n",
        "    - call /user/new\n",
        "    - create key for user\n",
        "    \"\"\"\n",
        "    global proxy_base_url\n",
        "    if not proxy_base_url.endswith(\"/\"):\n",
        "        proxy_base_url += \"/\"\n",
        "    url = proxy_base_url + \"user/new\"\n",
        "\n",
        "    # call /user/new\n",
        "    await client.post(\n",
        "        url=url,\n",
        "        headers={\"Authorization\": f\"Bearer {master_key}\"},\n",
        "        data=json.dumps({\n",
        "            \"user_id\": user_id,\n",
        "            \"user_alias\": user_name,\n",
        "            \"auto_create_key\": False,\n",
        "            # \"max_budget\": max_budget # 👈 [OPTIONAL] Sets max budget per user (if you want to set a max budget across keys)\n",
        "        })\n",
        "    )\n",
        "\n",
        "    # create key for user\n",
        "    return await create_key_with_alias(client=client, user_id=user_id, max_budget=max_budget)\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}