{
 "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\n",
    "import 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
}