{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# LiteLLM A121 Tutorial\n",
        "\n",
        "This walks through using A121 Jurassic models\n",
        "* j2-light\n",
        "* j2-mid\n",
        "* j2-ultra"
      ],
      "metadata": {
        "id": "LeFYo8iqcn5g"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GslPQFmaZsp-"
      },
      "outputs": [],
      "source": [
        "!pip install litellm"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from litellm import completion\n",
        "import os"
      ],
      "metadata": {
        "id": "P3cKiqURZx7P"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Set A121 Keys\n",
        "You can get a free key from https://studio.ai21.com/account/api-key"
      ],
      "metadata": {
        "id": "tmTvA1_GaNU4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "os.environ[\"AI21_API_KEY\"] = \"\""
      ],
      "metadata": {
        "id": "_xX8LmxAZ2vp"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# A121 Supported Models:\n",
        "https://studio.ai21.com/foundation-models"
      ],
      "metadata": {
        "id": "Fx5ZfJTLbF0A"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## J2-light Call"
      ],
      "metadata": {
        "id": "H0tl-0Z3bDaL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "messages = [{ \"content\": \"Hello, how are you?\",\"role\": \"user\"}]\n",
        "response = completion(model=\"j2-light\", messages=messages)\n",
        "response"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DZnApsJUZ_I2",
        "outputId": "b5707cbe-f67c-47f7-bac5-a7b8af1ba815"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<ModelResponse at 0x7b2c2902e610> JSON: {\n",
              "  \"choices\": [\n",
              "    {\n",
              "      \"finish_reason\": \"stop\",\n",
              "      \"index\": 0,\n",
              "      \"message\": {\n",
              "        \"content\": \" However, I have an important question to ask you\\nMy name is X, and I was wondering if you would be willing to help me.\",\n",
              "        \"role\": \"assistant\"\n",
              "      }\n",
              "    }\n",
              "  ],\n",
              "  \"created\": 1692761063.5189915,\n",
              "  \"model\": \"j2-light\",\n",
              "  \"usage\": {\n",
              "    \"prompt_tokens\": null,\n",
              "    \"completion_tokens\": null,\n",
              "    \"total_tokens\": null\n",
              "  }\n",
              "}"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# J2-Mid"
      ],
      "metadata": {
        "id": "wCcnrYnnbMQA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "messages = [{ \"content\": \"what model are you\",\"role\": \"user\"}]\n",
        "response = completion(model=\"j2-mid\", messages=messages)\n",
        "response"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-5Sxf4blaeEl",
        "outputId": "6264a5e8-16d6-44a3-e167-9e0c59b6dbc4"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<ModelResponse at 0x7b2c2902f6a0> JSON: {\n",
              "  \"choices\": [\n",
              "    {\n",
              "      \"finish_reason\": \"stop\",\n",
              "      \"index\": 0,\n",
              "      \"message\": {\n",
              "        \"content\": \"\\nplease choose the model from the list below\\nModel view in Tekla Structures\",\n",
              "        \"role\": \"assistant\"\n",
              "      }\n",
              "    }\n",
              "  ],\n",
              "  \"created\": 1692761140.0017524,\n",
              "  \"model\": \"j2-mid\",\n",
              "  \"usage\": {\n",
              "    \"prompt_tokens\": null,\n",
              "    \"completion_tokens\": null,\n",
              "    \"total_tokens\": null\n",
              "  }\n",
              "}"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# J2-Ultra"
      ],
      "metadata": {
        "id": "wDARpjxtbUcg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "messages = [{ \"content\": \"what model are you\",\"role\": \"user\"}]\n",
        "response = completion(model=\"j2-ultra\", messages=messages)\n",
        "response"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "i228xwsYbSYo",
        "outputId": "3765ac56-5a9b-442e-b357-2e346d02e1df"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<ModelResponse at 0x7b2c28fd4090> JSON: {\n",
              "  \"choices\": [\n",
              "    {\n",
              "      \"finish_reason\": \"stop\",\n",
              "      \"index\": 0,\n",
              "      \"message\": {\n",
              "        \"content\": \"\\nI am not a specific model, but I can provide information and assistance based on my training data. Please let me know if there is anything you\",\n",
              "        \"role\": \"assistant\"\n",
              "      }\n",
              "    }\n",
              "  ],\n",
              "  \"created\": 1692761157.8675153,\n",
              "  \"model\": \"j2-ultra\",\n",
              "  \"usage\": {\n",
              "    \"prompt_tokens\": null,\n",
              "    \"completion_tokens\": null,\n",
              "    \"total_tokens\": null\n",
              "  }\n",
              "}"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    }
  ]
}