diff --git a/docs/docs/integrations/llms/ai21.ipynb b/docs/docs/integrations/llms/ai21.ipynb
index 3adbd709874..2e22f85f11d 100644
--- a/docs/docs/integrations/llms/ai21.ipynb
+++ b/docs/docs/integrations/llms/ai21.ipynb
@@ -14,12 +14,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "02be122d-04e8-4ec6-84d1-f1d8961d6828",
    "metadata": {
     "tags": []
    },
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[33mWARNING: There was an error checking the latest version of pip.\u001b[0m\u001b[33m\n",
+      "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n"
+     ]
+    }
+   ],
    "source": [
     "# install the package:\n",
     "%pip install --upgrade --quiet  ai21"
@@ -27,20 +36,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 1,
    "id": "4229227e-6ca2-41ad-a3c3-5f29e3559091",
    "metadata": {
     "tags": []
    },
-   "outputs": [
-    {
-     "name": "stdin",
-     "output_type": "stream",
-     "text": [
-      " ········\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "# get AI21_API_KEY. Use https://studio.ai21.com/account/account\n",
     "\n",
@@ -51,21 +52,20 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 7,
    "id": "6fb585dd",
    "metadata": {
     "tags": []
    },
    "outputs": [],
    "source": [
-    "from langchain.chains import LLMChain\n",
-    "from langchain.prompts import PromptTemplate\n",
-    "from langchain_community.llms import AI21"
+    "from langchain_community.llms import AI21\n",
+    "from langchain_core.prompts import PromptTemplate"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 12,
    "id": "035dea0f",
    "metadata": {
     "tags": []
@@ -76,12 +76,12 @@
     "\n",
     "Answer: Let's think step by step.\"\"\"\n",
     "\n",
-    "prompt = PromptTemplate(template=template, input_variables=[\"question\"])"
+    "prompt = PromptTemplate.from_template(template)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 9,
    "id": "3f3458d9",
    "metadata": {
     "tags": []
@@ -93,19 +93,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 10,
    "id": "a641dbd9",
    "metadata": {
     "tags": []
    },
    "outputs": [],
    "source": [
-    "llm_chain = LLMChain(prompt=prompt, llm=llm)"
+    "llm_chain = prompt | llm"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 13,
    "id": "9f0b1960",
    "metadata": {
     "tags": []
@@ -114,10 +114,10 @@
     {
      "data": {
       "text/plain": [
-       "'\\n1. What year was Justin Bieber born?\\nJustin Bieber was born in 1994.\\n2. What team won the Super Bowl in 1994?\\nThe Dallas Cowboys won the Super Bowl in 1994.'"
+       "'\\nThe Super Bowl in the year Justin Beiber was born was in the year 1991.\\nThe Super Bowl in 1991 was won by the Washington Redskins.\\nFinal answer: Washington Redskins'"
       ]
      },
-     "execution_count": 12,
+     "execution_count": 13,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -125,7 +125,7 @@
    "source": [
     "question = \"What NFL team won the Super Bowl in the year Justin Beiber was born?\"\n",
     "\n",
-    "llm_chain.run(question)"
+    "llm_chain.invoke({\"question\": question})"
    ]
   },
   {
@@ -153,7 +153,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.3"
+   "version": "3.10.13"
   }
  },
  "nbformat": 4,