ai-gateway/resilient_loadbalancing_with_failure_mitigating_fallbacks.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Portkey-AI/portkey-cookbook/blob/main/ai-gateway/resilient_loadbalancing_with_failure_mitigating_fallbacks.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZcNFLOJCSLz8"
      },
      "source": [
        "# Setting up resilient Load balancers with failure-mitigating Fallbacks\n",
        "\n",
        "Companies often face challenges of scaling their services efficiently as the traffic to their applications grow - when you’re consuming APIs, the first point of failure is that if you hit the API too much, you can get rate limited. Loadbalancing is a proven way to scale usage horizontally without overburdening any one provider and thus staying within rate limits.\n",
        "\n",
        "For your AI app, rate limits are even more stringent, and if you start hitting the providers’ rate limits, there’s nothing you can do except wait to cool down and try again. With Portkey, we help you solve this very easily.\n",
        "\n",
        "This cookbook will teach you how to utilize Portkey to distribute traffic across multiple LLMs, ensuring that your loadbalancer is robust by setting up backups for requests. Additionally, you will learn how to load balance across OpenAI and Anthropic, leveraging the powerful Claude-3 models recently developed by Anthropic, with Azure serving as the fallback layer.\n",
        "\n",
        "<span style=\"text-decoration:underline;\">Prerequisites:</span>\n",
        "\n",
        "You should have the [Portkey API Key](https://portkey.ai/docs/api-reference/authentication#obtaining-your-api-key). Please sign up to obtain it. Additionally, you should have stored the OpenAI, Azure OpenAI, and Anthropic details in the [Portkey vault](https://portkey.ai/docs/product/ai-gateway-streamline-llm-integrations/virtual-keys)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iLIIavv5ShTm"
      },
      "source": [
        "## 1. Import the SDK and authenticate Portkey\n",
        "\n",
        "Start by installing the `portkey-ai` to your NodeJS project."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kcWetUoCSjo-",
        "outputId": "386bd3a6-4544-44c4-87c1-8077f864ae06"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Installing collected packages: pathspec, mypy-extensions, h11, mypy, httpcore, black, httpx, openai, portkey-ai\n",
            "Successfully installed black-23.7.0 h11-0.14.0 httpcore-1.0.5 httpx-0.27.0 mypy-1.9.0 mypy-extensions-1.0.0 openai-1.14.3 pathspec-0.12.1 portkey-ai-1.2.2\n"
          ]
        }
      ],
      "source": [
        "!pip install portkey-ai"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S1jiBF68StmV"
      },
      "source": [
        "Once installed, you can import it and instantiate it with the API key to your Portkey account."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "vejdxNwzSxul"
      },
      "outputs": [],
      "source": [
        "from portkey_ai import Portkey\n",
        "from google.colab import userdata\n",
        "\n",
        "PORTKEYAI_API_KEY=userdata.get('PORTKEY_API_KEY')\n",
        "\n",
        "portkey = Portkey(\n",
        "    api_key=PORTKEYAI_API_KEY\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rZ0ajEo2S-Ly"
      },
      "source": [
        "## 2. Create Configs: Loadbalance with Nested Fallbacks\n",
        "\n",
        "Portkey acts as AI gateway to all of your requests to LLMs. It follows the OpenAI SDK signature in all of it’s methods and interfaces making it easy to use and switch. Here is an example of an chat completions requests through Portkey."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BjnZ0MtZTTRC"
      },
      "outputs": [],
      "source": [
        "messages = [\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"content\": \"What are the 7 wonders of the world?\"\n",
        "    }\n",
        "]\n",
        "\n",
        "response = portkey.chat.completions.create(\n",
        "    messages = messages,\n",
        "    model = 'gpt-4'\n",
        ")\n",
        "\n",
        "print(response.choices[0].message.content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CogVqxkOTV69"
      },
      "source": [
        "The Portkey AI gateway can apply our desired behaviour to the requests to various LLMs. In a nutshell, our desired behaviour is the following:\n",
        "\n",
        "![](https://raw.githubusercontent.com/Saif-Shines/pk-cookbook/next/ai-gateway/images/resilient-loadbalancing-with-failure-mitigating-fallbacks/1-loadbalancing-with-fallbacks.png)\n",
        "\n",
        "Lucky for us, all of this can implemented by passing a configs allowing us to express what behavior to apply to every request through the Portkey AI gateway."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u1Yd2Zz2T_ec"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "\n",
        "config_data ={\n",
        "  \"strategy\":{\n",
        "    \"mode\":\"loadbalance\"\n",
        "  },\n",
        "  \"targets\":[\n",
        "    {\n",
        "      \"virtual_key\":\"ANTHROPIC_VIRTUAL_KEY\",\n",
        "      \"weight\":0.5,\n",
        "      \"override_params\":{\n",
        "        \"max_tokens\":200,\n",
        "        \"model\":\"claude-3-opus-20240229\"\n",
        "      }\n",
        "    },\n",
        "    {\n",
        "      \"strategy\":{\n",
        "        \"mode\":\"fallback\"\n",
        "      },\n",
        "      \"targets\":[\n",
        "        {\n",
        "          \"virtual_key\":\"OPENAI_VIRTUAL_KEY\"\n",
        "        },\n",
        "        {\n",
        "          \"virtual_key\":\"AZURE_OPENAI_VIRTUAL_KEY\"\n",
        "        }\n",
        "      ],\n",
        "      \"weight\":0.5\n",
        "    }\n",
        "  ]\n",
        "}\n",
        "\n",
        "portkey = Portkey(\n",
        "    api_key=PORTKEYAI_API_KEY,\n",
        "    config=json.dumps(config_data)\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iHHkc90EVLws"
      },
      "source": [
        "We apply the `loadbalance` strategy across _Anthropic and OpenAI._ `weight` describes the traffic should be split into 50/50 among both the LLM providers while `override_params` will help us override the defaults.\n",
        "\n",
        "Let’s take this a step further to apply a fallback mechanism for the requests from* OpenAI* to fallback to _Azure OpenAI_. This nested mechanism among the `targets` will ensure our app is reliable in the production in great confidence.\n",
        "\n",
        "See the documentation for Portkey [Fallbacks](https://portkey.ai/docs/product/ai-gateway-streamline-llm-integrations/fallbacks) and [Loadbalancing](https://portkey.ai/docs/product/ai-gateway-streamline-llm-integrations/load-balancing)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LBsIJUfbVPtc"
      },
      "source": [
        "## 3. Make a Request\n",
        "\n",
        "Now that the `config` ‘s are concrete and are passed as arguments when instantiating the Portkey client instance, all subsequent will acquire desired behavior auto-magically — No additional changes to the codebase."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VTENwNTLVT8l"
      },
      "outputs": [],
      "source": [
        "messages = [\n",
        "    {\n",
        "    \"role\": 'system',\n",
        "    \"content\": 'You are a very helpful assistant.'\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"content\": \"What are the 7 wonders of the world?\"\n",
        "    }\n",
        "]\n",
        "\n",
        "response = portkey.chat.completions.create(\n",
        "    messages = messages,\n",
        "    model = 'gpt-4'\n",
        ")\n",
        "\n",
        "print(response.choices[0].message.content) # The Seven Wonders of the Ancient World are:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "04pSEE3gVlhN"
      },
      "source": [
        "Next, we will examine how to identify load-balanced requests or those that have been executed as fallbacks."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_hhJZDAtVoc4"
      },
      "source": [
        "## 4. Trace the request from the logs\n",
        "\n",
        "It can be challenging to identify particular requests from the thousands that are received every day, similar to trying to find a needle in a haystack. However, Portkey offers a solution by enabling us to attach a desired trace ID. Here `request-loadbalance-fallback`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XBiVvC7lWCfD"
      },
      "outputs": [],
      "source": [
        "response = portkey.with_options(\n",
        "  trace_id=\"request-loadbalance-fallback\"\n",
        "  ).chat.completions.create(\n",
        "    messages = messages,\n",
        "    model = 'gpt-4'\n",
        ")\n",
        "\n",
        "print(response.choices[0].message.content)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NrhfKI5DWHPw"
      },
      "source": [
        "This trace ID can be used to filter requests from the Portkey Dashboard (>Logs) easily.\n",
        "\n",
        "![](https://raw.githubusercontent.com/Saif-Shines/pk-cookbook/next/ai-gateway/images/resilient-loadbalancing-with-failure-mitigating-fallbacks/2-logs-request-loadbalance-fallback.png)\n",
        "\n",
        "In addition to activating Loadbalance (icon), the logs provide essential observability information, including tokens, cost, and model.\n",
        "\n",
        "Are the configs growing and becoming harder to manage in the code? [Try creating them from Portkey UI](https://portkey.ai/docs/product/ai-gateway-streamline-llm-integrations/configs#creating-configs) and reference the configs ID in your code. It will make it significantly easier to maintain."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tWecz_GnWdoz"
      },
      "source": [
        "## 5. Advanced: Canary Testing\n",
        "\n",
        "Given there are new models coming every day and your app is in production — What is the best way to try the quality of those models? Canary Testing allows you to gradually roll out a change to a small subset of users before making it available to everyone.\n",
        "\n",
        "Consider this scenario: You have been using OpenAI as your LLM provider for a while now, but are considering trying an open-source Llama model for your app through Anyscale."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_A_MA9TtWoKF"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "\n",
        "OPENAI_VIRTUAL_KEY=\"\"\n",
        "ANYSCALE_VIRTUAL_KEY=\"\"\n",
        "\n",
        "config_data = {\n",
        "  \"strategy\": {\n",
        "    \"mode\": \"loadbalance\"\n",
        "  },\n",
        "  \"targets\": [\n",
        "    {\n",
        "      \"virtual_key\": OPENAI_VIRTUAL_KEY,\n",
        "      \"weight\": 0.9\n",
        "    },\n",
        "    {\n",
        "      \"virtual_key\": ANYSCALE_VIRTUAL_KEY,\n",
        "      \"weight\": 0.1,\n",
        "      \"override_params\": {\n",
        "        \"model\": \"meta-llama/Llama-2-70b-chat-hf\"\n",
        "      }\n",
        "    }\n",
        "  ]\n",
        "}\n",
        "\n",
        "portkey = Portkey(\n",
        "    api_key=PORTKEYAI_API_KEY,\n",
        "    config=json.dumps(config_data)\n",
        ")\n",
        "\n",
        "messages = [\n",
        "    {\n",
        "      \"role\": 'system',\n",
        "      \"content\": 'You are a very helpful assistant.'\n",
        "    },\n",
        "    {\n",
        "      \"role\": \"user\",\n",
        "      \"content\": \"What are the 7 wonders of the world?\"\n",
        "    }\n",
        "]\n",
        "\n",
        "response = portkey.with_options(\n",
        "  trace_id='canary-testing'\n",
        "  ).chat.completions.create(\n",
        "    messages = messages,\n",
        "    model = 'gpt-3.5-turbo'\n",
        ")\n",
        "\n",
        "print(response.choices[0].message.content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eJmDj4pxW648"
      },
      "source": [
        "\n",
        "The `weight` , indication of traffic is split to have 10% of your user-base are served from Anyscale’s Llama models. Now, you are all set up to get feedback and observe the performance of your app and release increasingly to larger userbase."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f41eebD3W9XK"
      },
      "source": [
        "## Considerations\n",
        "\n",
        "You can implement production-grade Loadbalancing and nested fallback mechanisms with just a few lines of code. While you are equipped with all the tools for your next GenAI app, here are a few considerations:\n",
        "\n",
        "- Every request has to adhere to the LLM provider’s requirements for it to be successful. For instance, `max_tokens` is required for Anthropic and not for OpenAI.\n",
        "- While loadbalance helps reduce the load on one LLM - it is recommended to pair it with a Fallback strategy to ensure that your app stays reliable\n",
        "- On Portkey, you can also pass the loadbalance weight as 0 - this will essentially stop routing requests to that target and you can amp it up when required\n",
        "- Loadbalance has no target limits as such, so you can potentially add multiple account details from one provider and effectively multiply your available rate limits\n",
        "- Loadbalance does not alter the outputs or the latency of the requests in any way\n",
        "\n",
        "Happy Coding!"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "authorship_tag": "ABX9TyNreh2pqZM2v+EntI5rVr96",
      "include_colab_link": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}