From d1cb372b95cd311cb5288036adc6f76a98742e2a Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Tue, 26 Aug 2025 21:09:15 +0800
Subject: [PATCH] fix notebook

---
 cookbook/pageindex_RAG_simple.ipynb | 43 ++++++++++++++++++++++++-----
 1 file changed, 36 insertions(+), 7 deletions(-)

diff --git a/cookbook/pageindex_RAG_simple.ipynb b/cookbook/pageindex_RAG_simple.ipynb
index a56d15f..0c6be0a 100644
--- a/cookbook/pageindex_RAG_simple.ipynb
+++ b/cookbook/pageindex_RAG_simple.ipynb
@@ -94,7 +94,7 @@
         "id": "edTfrizMFK4c"
       },
       "source": [
-        "#### 0.1 Install dependencies"
+        "#### 0.1 Install PageIndex"
       ]
     },
     {
@@ -106,7 +106,7 @@
       },
       "outputs": [],
       "source": [
-        "%pip install -q --upgrade pageindex openai"
+        "%pip install -q --upgrade pageindex"
       ]
     },
     {
@@ -115,7 +115,7 @@
         "id": "WVEWzPKGcG1M"
       },
       "source": [
-        "#### 0.2 Setup environment"
+        "#### 0.2 Setup PageIndex"
       ]
     },
     {
@@ -126,15 +126,40 @@
       },
       "outputs": [],
       "source": [
-        "import json, os, requests\n",
         "from pageindex import PageIndexClient\n",
         "import pageindex.utils as utils\n",
         "\n",
         "# Get your PageIndex API key from https://dash.pageindex.ai/api-keys\n",
         "PAGEINDEX_API_KEY = \"YOUR_PAGEINDEX_API_KEY\"\n",
+        "pi_client = PageIndexClient(api_key=PAGEINDEX_API_KEY)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "#### 0.3 Setup LLM\n",
+        "\n",
+        "Choose your preferred LLM for reasoning-based retrieval. In this example, we use OpenAI’s GPT-4.1."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import openai\n",
         "OPENAI_API_KEY = \"YOUR_OPENAI_API_KEY\"\n",
         "\n",
-        "pi_client = PageIndexClient(api_key=PAGEINDEX_API_KEY)"
+        "async def call_llm(prompt, model=\"gpt-4.1\", temperature=0):\n",
+        "    client = openai.AsyncOpenAI(api_key=OPENAI_API_KEY)\n",
+        "    response = await client.chat.completions.create(\n",
+        "        model=model,\n",
+        "        messages=[{\"role\": \"user\", \"content\": prompt}],\n",
+        "        temperature=temperature\n",
+        "    )\n",
+        "    return response.choices[0].message.content.strip()"
       ]
     },
     {
@@ -176,6 +201,8 @@
         }
       ],
       "source": [
+        "import os, requests\n",
+        "\n",
         "# You can also use our GitHub repo to generate PageIndex tree\n",
         "# https://github.com/VectifyAI/PageIndex\n",
         "\n",
@@ -321,6 +348,8 @@
       },
       "outputs": [],
       "source": [
+        "import json\n",
+        "\n",
         "query = \"What are the conclusions in this document?\"\n",
         "\n",
         "tree_without_text = utils.remove_fields(tree.copy(), fields=['text'])\n",
@@ -343,7 +372,7 @@
         "Directly return the final JSON structure. Do not output anything else.\n",
         "\"\"\"\n",
         "\n",
-        "tree_search_result = await utils.call_llm(search_prompt, api_key=OPENAI_API_KEY)"
+        "tree_search_result = await call_llm(search_prompt)"
       ]
     },
     {
@@ -508,7 +537,7 @@
         "\"\"\"\n",
         "\n",
         "print('Generated Answer:\\n')\n",
-        "answer = await utils.call_llm(answer_prompt, api_key=OPENAI_API_KEY)\n",
+        "answer = await call_llm(answer_prompt)\n",
         "utils.print_wrapped(answer)"
       ]
     },