feat: enhancements and updates for portkey (run-llama#7669)

CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ## Unreleased
 
+### New Features
+- Simplified portkey LLM interface (#7669)
+
 ### Bug Fixes / Nits
 - Avoid `NotImplementedError` for async langchain embeddings (#7668)
 

docs/examples/llm/portkey.ipynb

@@ -33,19 +33,14 @@
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
+   "cell_type": "markdown",
    "metadata": {
-    "id": "5Z933R9wuZ4z"
+    "id": "4Q8L0IDau60f"
    },
-   "outputs": [],
    "source": [
-    "# Installing the Rubeus AI gateway developed by the Portkey team\n",
-    "!pip install rubeus\n",
+    "#### **Step 1: Get your Portkey API key**\n",
     "\n",
-    "# Importing necessary libraries and modules\n",
-    "from llama_index.llms import Portkey, ChatMessage\n",
-    "from rubeus import LLMBase"
+    "Log into [Portkey here](https://app.portkey.ai/), then click on the profile icon on top right and \"Copy API Key\". Let's also set OpenAI & Anthropic API keys."
    ]
   },
   {
@@ -60,29 +55,35 @@
    ]
   },
   {
-   "cell_type": "markdown",
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {
-    "id": "4Q8L0IDau60f"
+    "id": "Vd7elJZHu_jF"
    },
+   "outputs": [],
    "source": [
-    "#### **Step 1: Get your Portkey API key**\n",
+    "# !pip install portkey-ai -U\n",
     "\n",
-    "Log into [Portkey here](https://app.portkey.ai/), then click on the profile icon on top right and \"Copy API Key\". Let's also set OpenAI & Anthropic API keys."
+    "# Set the portkey api key as environment variable.\n",
+    "import os\n",
+    "\n",
+    "os.environ[\"PORTKEY_API_KEY\"] = \"PORTKEY_API_KEY\"\n",
+    "os.environ[\"OPENAI_API_KEY\"] = \"\"\n",
+    "os.environ[\"ANTHROPIC_API_KEY\"] = \"\""
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "id": "Vd7elJZHu_jF"
+    "id": "5Z933R9wuZ4z"
    },
    "outputs": [],
    "source": [
-    "import os\n",
-    "\n",
-    "os.environ[\"PORTKEY_API_KEY\"] = \"\"\n",
-    "os.environ[\"OPENAI_API_KEY\"] = \"\"\n",
-    "os.environ[\"ANTHROPIC_API_KEY\"] = \"\""
+    "# Installing the Portkey's AI gateway SDK developed by the Portkey team\n",
+    "# Importing necessary libraries and modules\n",
+    "from llama_index.llms import Portkey, ChatMessage\n",
+    "import portkey as pk"
    ]
   },
   {
@@ -118,21 +119,8 @@
    },
    "outputs": [],
    "source": [
-    "metadata = {\n",
-    "    \"_environment\": \"production\",\n",
-    "    \"_prompt\": \"test\",\n",
-    "    \"_user\": \"user\",\n",
-    "    \"_organisation\": \"acme\",\n",
-    "}\n",
-    "\n",
-    "pk_llm = Portkey(\n",
+    "portkey_client = Portkey(\n",
     "    mode=\"single\",\n",
-    "    cache_status=\"semantic\",\n",
-    "    cache_force_refresh=True,\n",
-    "    cache_age=1000,\n",
-    "    trace_id=\"portkey_llamaindex\",\n",
-    "    retry=5,\n",
-    "    metadata=metadata,\n",
     ")\n",
     "\n",
     "# Since we have defined the Portkey API Key with os.environ, we do not need to set api_key again here"
@@ -146,7 +134,7 @@
    "source": [
     "#### **Step 3: Constructing the LLM**\n",
     "\n",
-    "With the Portkey integration, constructing an LLM is simplified. Use the `LLMBase` function for all providers, with the exact same keys you're accustomed to in your OpenAI or Anthropic constructors. The only new key is `weight`, essential for the load balancing feature."
+    "With the Portkey integration, constructing an LLM is simplified. Use the `LLMOptions` function for all providers, with the exact same keys you're accustomed to in your OpenAI or Anthropic constructors. The only new key is `weight`, essential for the load balancing feature."
    ]
   },
   {
@@ -157,7 +145,9 @@
    },
    "outputs": [],
    "source": [
-    "openai_llm = LLMBase(provider=\"openai\", model=\"gpt-4\")"
+    "openai_llm = pk.LLMOptions(\n",
+    "    provider=\"openai\", model=\"gpt-4\", virtual_key=\"open-ai-key-66ah788\"\n",
+    ")"
    ]
   },
   {
@@ -166,7 +156,7 @@
     "id": "7foKVwDdwBgG"
    },
    "source": [
-    "The above code illustrates how to utilize the `LLMBase` function to set up an LLM with the OpenAI provider and the GPT-4 model. This same function can be used for other providers as well, making the integration process streamlined and consistent across various providers."
+    "The above code illustrates how to utilize the `LLMOptions` function to set up an LLM with the OpenAI provider and the GPT-4 model. This same function can be used for other providers as well, making the integration process streamlined and consistent across various providers."
    ]
   },
   {
@@ -177,7 +167,7 @@
    "source": [
     "#### **Step 4: Activate the Portkey LLM**\n",
     "\n",
-    "Once you've constructed the LLM using the `LLMBase` function, the next step is to activate it with Portkey. This step is essential to ensure that all the Portkey features are available for your LLM."
+    "Once you've constructed the LLM using the `LLMOptions` function, the next step is to activate it with Portkey. This step is essential to ensure that all the Portkey features are available for your LLM."
    ]
   },
   {
@@ -188,7 +178,7 @@
    },
    "outputs": [],
    "source": [
-    "pk_llm.add_llms(openai_llm)"
+    "portkey_client.add_llms(openai_llm)"
    ]
   },
   {
@@ -224,7 +214,7 @@
     "    ChatMessage(role=\"user\", content=\"What can you do?\"),\n",
     "]\n",
     "print(\"Testing Portkey Llamaindex integration:\")\n",
-    "response = pk_llm.chat(messages)\n",
+    "response = portkey_client.chat(messages)\n",
     "print(response)"
    ]
   },
@@ -238,17 +228,17 @@
     "\n",
     "Congratulations! You've successfully set up and tested the Portkey integration with Llamaindex. To recap:\n",
     "\n",
-    "1. pip install rubeus\n",
+    "1. pip install portkey-ai\n",
     "2. Import Portkey from llama_index.llms.\n",
     "3. Grab your Portkey API Key from [here](https://app.portkey.ai/).\n",
-    "4. Construct your Portkey LLM with `pk_llm=Portkey(mode=\"fallback\")` and any other Portkey features you want\n",
-    "5. Construct your provider LLM with opneai_llm=PortkeyBase(provider=\"openai\",model=\"gpt-4\")\n",
-    "6. Add the provider LLM to Portkey LLM with `pk_llm.add_llms(openai_llm)`\n",
-    "7. Call the Portkey LLM regularly like you would any other LLM, with `pk_llm.chat(messages)`\n",
+    "4. Construct your Portkey client with `portkey_client=Portkey(mode=\"fallback\")` and any other Portkey modes you want\n",
+    "5. Construct your provider LLM with `opneai_llm=PortkeyBase(provider=\"openai\",model=\"gpt-4\")`\n",
+    "6. Add the provider LLM to Portkey LLM with `portkey_client.add_llms(openai_llm)`\n",
+    "7. Call the Portkey methods regularly like you would any other LLM, with `portkey_client.chat(messages)`\n",
     "\n",
     "Here's the guide to all the functions and their params:\n",
     "- [Portkey LLM Constructor](#step-2-add-all-the-portkey-features-you-want-as-illustrated-below-by-calling-the-portkey-class)\n",
-    "- [LLMBase Constructor](https://github.com/Portkey-AI/rubeus-python-sdk/blob/4cf3e17b847225123e92f8e8467b41d082186d60/rubeus/api_resources/utils.py#L179)\n",
+    "- [LLMOptions Constructor](https://github.com/Portkey-AI/rubeus-python-sdk/blob/4cf3e17b847225123e92f8e8467b41d082186d60/rubeus/api_resources/utils.py#L179)\n",
     "- [List of Portkey + Llamaindex Features](#portkeys-integration-with-llamaindex-adds-the-following-production-capabilities-to-your-apps-out-of-the-box)\n"
    ]
   },
@@ -276,15 +266,26 @@
    },
    "outputs": [],
    "source": [
-    "pk_llm = Portkey(mode=\"fallback\", retry=5)\n",
+    "portkey_client = Portkey(mode=\"fallback\")\n",
+    "messages = [\n",
+    "    ChatMessage(role=\"system\", content=\"You are a helpful assistant\"),\n",
+    "    ChatMessage(role=\"user\", content=\"What can you do?\"),\n",
+    "]\n",
     "\n",
-    "llm1 = LLMBase(provider=\"openai\", model=\"gpt-4\")\n",
-    "llm2 = LLMBase(provider=\"openai\", model=\"gpt-3.5-turbo\")\n",
+    "llm1 = pk.LLMOptions(\n",
+    "    provider=\"openai\",\n",
+    "    model=\"gpt-4\",\n",
+    "    retry_settings={\"on_status_codes\": [429, 500], \"attempts\": 2},\n",
+    "    virtual_key=\"open-ai-key-66ah788\",\n",
+    ")\n",
+    "llm2 = pk.LLMOptions(\n",
+    "    provider=\"openai\", model=\"gpt-3.5-turbo\", virtual_key=\"open-ai-key-66ah788\"\n",
+    ")\n",
     "\n",
-    "pk_llm.add_llms(llm_params=[llm1, llm2])\n",
+    "portkey_client.add_llms(llm_params=[llm1])\n",
     "\n",
     "print(\"Testing Fallback & Retry functionality:\")\n",
-    "response = pk_llm.chat(messages)\n",
+    "response = portkey_client.chat(messages)\n",
     "print(response)"
    ]
   },
@@ -314,15 +315,26 @@
    },
    "outputs": [],
    "source": [
-    "pk_llm = Portkey(mode=\"loadbalance\")\n",
+    "portkey_client = Portkey(mode=\"ab_test\")\n",
+    "messages = [\n",
+    "    ChatMessage(role=\"system\", content=\"You are a helpful assistant\"),\n",
+    "    ChatMessage(role=\"user\", content=\"What can you do?\"),\n",
+    "]\n",
     "\n",
-    "llm1 = LLMBase(provider=\"openai\", model=\"gpt-4\", weight=0.2)\n",
-    "llm2 = LLMBase(provider=\"openai\", model=\"gpt-3.5-turbo\", weight=0.8)\n",
+    "llm1 = pk.LLMOptions(\n",
+    "    provider=\"openai\", virtual_key=\"open-ai-key-66ah788\", model=\"gpt-4\", weight=0.2\n",
+    ")\n",
+    "llm2 = pk.LLMOptions(\n",
+    "    provider=\"openai\",\n",
+    "    virtual_key=\"open-ai-key-66ah788\",\n",
+    "    model=\"gpt-3.5-turbo\",\n",
+    "    weight=0.8,\n",
+    ")\n",
     "\n",
-    "pk_llm.add_llms(llm_params=[llm1, llm2])\n",
+    "portkey_client.add_llms(llm_params=[llm1, llm2])\n",
     "\n",
     "print(\"Testing Loadbalance functionality:\")\n",
-    "response = pk_llm.chat(messages)\n",
+    "response = portkey_client.chat(messages)\n",
     "print(response)"
    ]
   },
@@ -349,8 +361,15 @@
    "source": [
     "import time\n",
     "\n",
-    "pk_llm = Portkey(mode=\"single\", cache_status=\"semantic\")\n",
-    "pk_llm.add_llms(llm1)\n",
+    "portkey_client = Portkey(mode=\"single\")\n",
+    "\n",
+    "openai_llm = pk.LLMOptions(\n",
+    "    provider=\"openai\",\n",
+    "    virtual_key=\"open-ai-key-66ah788\",\n",
+    "    model=\"gpt-3.5-turbo\",\n",
+    "    cache_status=\"semantic\",\n",
+    ")\n",
+    "portkey_client.add_llms(openai_llm)\n",
     "\n",
     "current_messages = [\n",
     "    ChatMessage(role=\"system\", content=\"You are a helpful assistant\"),\n",
@@ -360,7 +379,7 @@
     "print(\"Testing Portkey Semantic Cache:\")\n",
     "\n",
     "start = time.time()\n",
-    "response = pk_llm.chat(current_messages)\n",
+    "response = portkey_client.chat(current_messages)\n",
     "end = time.time() - start\n",
     "\n",
     "print(response)\n",
@@ -377,7 +396,7 @@
     "print(\"Testing Portkey Semantic Cache:\")\n",
     "\n",
     "start = time.time()\n",
-    "response = pk_llm.chat(new_messages)\n",
+    "response = portkey_client.chat(new_messages)\n",
     "end = time.time() - start\n",
     "\n",
     "print(response)\n",
@@ -408,8 +427,13 @@
    },
    "outputs": [],
    "source": [
-    "pk_llm = Portkey(\n",
-    "    mode=\"single\", cache_status=\"semantic\", cache_age=1000, cache_force_refresh=True\n",
+    "# Setting the cache status as `semantic` and cache_age as 60s.\n",
+    "openai_llm = pk.LLMOptions(\n",
+    "    provider=\"openai\",\n",
+    "    virtual_key=\"open-ai-key-66ah788\",\n",
+    "    model=\"gpt-3.5-turbo\",\n",
+    "    cache_force_refresh=True,\n",
+    "    cache_age=50,\n",
     ")"
    ]
   },
@@ -441,14 +465,11 @@
     "    \"_organisation\": \"acme\",\n",
     "}\n",
     "\n",
-    "pk_llm = Portkey(\n",
-    "    mode=\"single\",\n",
-    "    trace_id=\"portkey_llamaindex_test\",\n",
-    ")\n",
-    "pk_llm.add_llms(llm1)\n",
+    "portkey_client = Portkey(mode=\"single\")\n",
+    "portkey_client.add_llms(openai_llm)\n",
     "\n",
     "print(\"Testing Observability functionality:\")\n",
-    "response = pk_llm.chat(messages)\n",
+    "response = portkey_client.chat(messages)\n",
     "print(response)"
    ]
   },
@@ -458,14 +479,14 @@
     "id": "KeOQL-9uxxQz"
    },
    "source": [
-    "#### **AI Gateway with Rubeus**\n",
+    "#### **AI Gateway with Portkey**\n",
     "\n",
-    "Rubeus is an open-source AI gateway that powers features like load balancing and fallbacks in Portkey. It acts as an intermediary, ensuring that your requests are processed optimally. One of the advantages of using Rubeus is its flexibility. You can easily customize its behavior, redirect requests to different providers, or even bypass logging to Portkey.\n",
+    "Portkey is an open-source AI gateway that powers features like load balancing and fallbacks. It acts as an intermediary, ensuring that your requests are processed optimally. One of the advantages of using Portkey is its flexibility. You can easily customize its behavior, redirect requests to different providers, or even bypass logging to Portkey.\n",
     "\n",
-    "Here's an example of customizing the behavior with Rubeus:\n",
+    "Here's an example of customizing the behavior with Portkey:\n",
     "\n",
     "```py\n",
-    "pk_llm.base_url=None\n",
+    "portkey_client.base_url=None\n",
     "```"
    ]
   },
@@ -550,7 +571,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.2"
+   "version": "3.11.5"
   }
  },
  "nbformat": 4,

examples/portkey_demo/fallback_demo.py

@@ -1,7 +1,7 @@
 import os
 from llama_index.llms import Portkey
 from llama_index.llms import ChatMessage  # We'll use this later
-from rubeus import LLMBase
+from portkey import LLMOptions
 
 os.environ["PORTKEY_API_KEY"] = ""
 os.environ["OPENAI_API_KEY"] = ""
@@ -16,19 +16,32 @@
 }
 
 # Define the Portkey interface.
-pk_llm = Portkey(mode="fallback", trace_id="portkey_llamaindex", metadata=metadata)
+pk_client = Portkey(mode="fallback")
 
 # Define the provider spec in the LLMBase spec. Customise the LLMs as per requirement.
-openai_llm = LLMBase(provider="openai", model="gpt-4")
-anthropic_llm = LLMBase(provider="openai", model="claude-2", max_tokens=256)
-
-pk_llm.add_llms([openai_llm, anthropic_llm])
+openai_llm = LLMOptions(
+    provider="openai",
+    model="gpt-4",
+    trace_id="portkey_llamaindex",
+    metadata=metadata,
+    virtual_key="open-ai-key-66ah788",
+)
+anthropic_llm = LLMOptions(
+    provider="openai",
+    model="claude-2",
+    max_tokens=256,
+    trace_id="portkey_llamaindex",
+    metadata=metadata,
+    virtual_key="anthropic-key-351feb",
+)
+
+pk_client.add_llms([openai_llm, anthropic_llm])
 
 messages = [
     ChatMessage(role="system", content="You are a helpful assistant"),
     ChatMessage(role="user", content="What can you do?"),
 ]
 
 print("Testing Portkey Llamaindex integration:")
-response = pk_llm.chat(messages)
+response = pk_client.chat(messages)
 print(response)