diff --git a/08-building-search-applications/python/aoai-solution.ipynb b/08-building-search-applications/python/aoai-solution.ipynb
index 1fc8a210a..ea9b5da82 100644
--- a/08-building-search-applications/python/aoai-solution.ipynb
+++ b/08-building-search-applications/python/aoai-solution.ipynb
@@ -22,7 +22,7 @@
     "load_dotenv()\n",
     "\n",
     "client = AzureOpenAI(\n",
-    "  api_key=os.environ['AZURE_OPENAI_KEY'],  # this is also the default, it can be omitted\n",
+    "  api_key=os.environ['AZURE_OPENAI_API_KEY'],  # this is also the default, it can be omitted\n",
     "  api_version = \"2023-05-15\"\n",
     "  )\n",
     "\n",
@@ -71,6 +71,10 @@
    "outputs": [],
    "source": [
     "def cosine_similarity(a, b):\n",
+    "    if len(a) > len(b):\n",
+    "        b = np.pad(b, (0, len(a) - len(b)), 'constant')\n",
+    "    elif len(b) > len(a):\n",
+    "        a = np.pad(a, (0, len(b) - len(a)), 'constant')\n",
     "    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))\n",
     "\n",
     "def get_videos(\n",